Digispark-Attiny85-payloads/duck2spark.py

#!/usr/bin/env python3

import getopt
import os
import sys


def generate_source(
    payload,
    init_delay=2500,
    loop_count=-1,
    loop_delay=5000,
    blink=True,
):
    head = '''/*
* Sketch generated by duck2spark from Marcus Mengs aka MaMe82
*
*/
#include "DigiKeyboard.h"

'''

    init = '''
void setup()
{
    // initialize the digital pin as an output.
    pinMode(0, OUTPUT); // LED on Model B
    pinMode(1, OUTPUT); // LED on Model A
    DigiKeyboard.delay(%d); // wait %d milliseconds before first run,
                            // to give target time to initialize
}

void loop()
{
''' % (init_delay, init_delay)

    body = '''
    // should code be runned in this loop?
    if (i != 0) {
        DigiKeyboard.sendKeyStroke(0);

        // parse raw duckencoder script
        for (int i = 0; i < DUCK_LEN; i += 2)
        {
            uint8_t key = pgm_read_word_near(duckraw + i);
            uint8_t mod = pgm_read_word_near(duckraw + i + 1);

            if (key == 0) // delay (a delay > 255 is split into a sequence of delays)
            {
                DigiKeyboard.delay(mod);
            }
            else
            {
                DigiKeyboard.sendKeyStroke(key, mod);
            }
        }

        i--;
        DigiKeyboard.delay(%d); // wait %d milliseconds before next loop iteration
    }
    else if (blink)
    {
        digitalWrite(0, HIGH);   // turn the LED on
        digitalWrite(1, HIGH);
        delay(100);

        digitalWrite(0, LOW);    // turn the LED off
        digitalWrite(1, LOW);
        delay(100);
    }
''' % (loop_delay, loop_delay)

    tail = '''}
'''

    # Generate the payload declaration in FLASH memory.
    payload_len = len(payload)

    declare = (
        "#define DUCK_LEN "
        + str(payload_len)
        + "\nconst PROGMEM uint8_t duckraw [DUCK_LEN] = {\n\t"
    )

    # In Python 3, payload is bytes, so payload[c] is already an int.
    for c in range(payload_len - 1):
        declare += f"{hex(payload[c])}, "

    if payload_len > 0:
        declare += f"{hex(payload[-1])}\n"

    declare += (
        "};\n"
        f"int i = {loop_count}; "
        "// how many times the payload should run (-1 for endless loop)\n"
    )

    if blink:
        declare += "bool blink=true;\n"
    else:
        declare += "bool blink=false;\n"

    return head + declare + init + body + tail


def usage():
    usagescr = '''MaMe82 duck2spark 1.0
=====================

Converts payload created by DuckEncoder to sourcefile for DigiSpark Sketch

Usage:
    python3 duck2spark.py -i [file ..]
        build Sketch from specified RubberDucky payload file

    python3 duck2spark.py -i [file ..] -o [file ..]
        save Sketch source to specified output file

Arguments:
   -i [file ..]     Input File (Payload encoded with DuckEncoder)
   -o [file ..]     Output File for Sketch, if omitted stdout is used
   -l <count>       Loop count (1=single run (default), -1=endless run,
                    3=3 runs etc.)
   -f <millis>      Delay in milliseconds before initial payload run
                    (default 1000)
   -r <millis>      Delay in milliseconds between loop runs
                    (default 5000)
   -n                Don't blink status LED after finish of payload execution

Remark:
    In order to use DEAD KEYS (e.g. ^ and ` on German keyboard layout)
    a SPACE should be appended in the ducky script
    (e.g. "STRING ^ working deadkey").
'''
    print(usagescr)


def main(argv):
    ifile = ""
    ofile = None
    payload = None

    loop_count = 1
    blink = True
    init_delay = 1000
    loop_delay = 5000

    try:
        opts, args = getopt.getopt(
            argv,
            "hi:o:l:nf:r:",
            [
                "help",
                "input=",
                "output=",
                "loopcount=",
                "noblink",
                "initdelay=",
                "repeatdelay=",
            ],
        )
    except getopt.GetoptError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        usage()
        sys.exit(2)

    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit(0)

        elif opt in ("-i", "--input"):
            ifile = arg

            if not os.path.isfile(ifile) or not os.access(ifile, os.R_OK):
                print(
                    f"Input file {ifile} doesn't exist or isn't readable",
                    file=sys.stderr,
                )
                sys.exit(2)

            with open(ifile, "rb") as f:
                payload = f.read()

        elif opt in ("-o", "--output"):
            ofile = arg

        elif opt in ("-l", "--loopcount"):
            try:
                loop_count = int(arg)
            except ValueError:
                print(
                    f"Invalid loop count: {arg}",
                    file=sys.stderr,
                )
                sys.exit(2)

        elif opt in ("-f", "--initdelay"):
            try:
                init_delay = int(arg)
            except ValueError:
                print(
                    f"Invalid initial delay: {arg}",
                    file=sys.stderr,
                )
                sys.exit(2)

        elif opt in ("-r", "--repeatdelay"):
            try:
                loop_delay = int(arg)
            except ValueError:
                print(
                    f"Invalid repeat delay: {arg}",
                    file=sys.stderr,
                )
                sys.exit(2)

        elif opt in ("-n", "--noblink"):
            blink = False

    if payload is None:
        print(
            "You have to provide a payload generated by DuckEncoder "
            "(-i option)",
            file=sys.stderr,
        )
        sys.exit(2)

    # Generate source code for Sketch.
    result = generate_source(
        payload,
        init_delay=init_delay,
        loop_count=loop_count,
        loop_delay=loop_delay,
        blink=blink,
    )

    if ofile is None:
        # Print to stdout.
        print(result, end="")
    else:
        # Write to output file.
        with open(ofile, "w", encoding="utf-8") as f:
            f.write(result)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        usage()
        sys.exit(1)

    main(sys.argv[1:])