#!/usr/bin/env python3

import getopt
import os
import sys
import time


class DuckEncoder:
    @staticmethod
    def readResource(filename):
        result_dict = {}

        with open(filename, "r", encoding="utf-8") as f:
            for line in f:
                # Remove comments.
                line = line.split("//", 1)[0]

                # Remove whitespace / line breaks.
                line = line.strip()

                # Skip empty lines.
                if not line:
                    continue

                key, val = line.split("=", 1)
                result_dict[key.strip()] = val.strip()

        return result_dict

    @staticmethod
    def parseScriptLine(line, keyProp, langProp):
        result = b""

        # Split line into command and arguments.
        cmd, _, args = line.partition(" ")
        cmd = cmd.strip()
        args = args.strip()

        # DELAY
        if cmd == "DELAY":
            delay = int(args)
            result = DuckEncoder.delay2USBBytes(delay)

        # STRING
        elif cmd == "STRING":
            if not args:
                return b""

            for c in args:
                keydata = DuckEncoder.ASCIIChar2USBBytes(
                    c, keyProp, langProp
                )
                if keydata:
                    result += keydata

        # STRING_DELAY
        elif cmd == "STRING_DELAY":
            if not args:
                return b""

            delay_str, chars = args.split(" ", 1)

            delay = int(delay_str.strip())
            delaystr = DuckEncoder.delay2USBBytes(delay)

            for c in chars.strip():
                keydata = DuckEncoder.ASCIIChar2USBBytes(
                    c, keyProp, langProp
                )
                if keydata:
                    result += keydata + delaystr

        # CONTROL / CTRL
        elif cmd in ("CONTROL", "CTRL"):
            if args:
                result = (
                    DuckEncoder.keyInstr2USBBytes(
                        args, keyProp, langProp
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_CTRL", keyProp, langProp
                    )
                )
            else:
                result = (
                    DuckEncoder.prop2USBByte(
                        "KEY_LEFT_CTRL", keyProp, langProp
                    )
                    + b"\x00"
                )

        # ALT
        elif cmd == "ALT":
            if args:
                result = (
                    DuckEncoder.keyInstr2USBBytes(
                        args, keyProp, langProp
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_ALT", keyProp, langProp
                    )
                )
            else:
                result = (
                    DuckEncoder.prop2USBByte(
                        "KEY_LEFT_ALT", keyProp, langProp
                    )
                    + b"\x00"
                )

        # SHIFT
        elif cmd == "SHIFT":
            if args:
                result = (
                    DuckEncoder.keyInstr2USBBytes(
                        args, keyProp, langProp
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_SHIFT", keyProp, langProp
                    )
                )
            else:
                result = (
                    DuckEncoder.prop2USBByte(
                        "KEY_LEFT_SHIFT", keyProp, langProp
                    )
                    + b"\x00"
                )

        # CTRL-ALT
        elif cmd == "CTRL-ALT":
            if args:
                key = DuckEncoder.keyInstr2USBBytes(
                    args, keyProp, langProp
                )

                ctrl = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_CTRL", keyProp, langProp
                )
                alt = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_ALT", keyProp, langProp
                )

                result = key + bytes([ctrl[0] | alt[0]])
            else:
                return b""

        # CTRL-SHIFT
        elif cmd == "CTRL-SHIFT":
            if args:
                key = DuckEncoder.keyInstr2USBBytes(
                    args, keyProp, langProp
                )

                ctrl = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_CTRL", keyProp, langProp
                )
                shift = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_SHIFT", keyProp, langProp
                )

                result = key + bytes([ctrl[0] | shift[0]])
            else:
                return b""

        # COMMAND-OPTION
        elif cmd == "COMMAND-OPTION":
            if args:
                key = DuckEncoder.keyInstr2USBBytes(
                    args, keyProp, langProp
                )

                command = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_LEFT_GUI", keyProp, langProp
                )
                alt = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_ALT", keyProp, langProp
                )

                result = key + bytes([command[0] | alt[0]])
            else:
                return b""

        # ALT-SHIFT
        elif cmd == "ALT-SHIFT":
            if args:
                key = DuckEncoder.keyInstr2USBBytes(
                    args, keyProp, langProp
                )

                alt = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_LEFT_ALT", keyProp, langProp
                )
                shift = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_SHIFT", keyProp, langProp
                )

                result = key + bytes([alt[0] | shift[0]])
            else:
                alt = DuckEncoder.prop2USBByte(
                    "KEY_LEFT_ALT", keyProp, langProp
                )

                shift = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_LEFT_ALT", keyProp, langProp
                )

                result = alt + bytes([alt[0] | shift[0]])

        # ALT-TAB
        elif cmd == "ALT-TAB":
            if args:
                return b""
            else:
                key = DuckEncoder.prop2USBByte(
                    "KEY_TAB", keyProp, langProp
                )
                alt = DuckEncoder.prop2USBByte(
                    "MODIFIERKEY_LEFT_ALT", keyProp, langProp
                )

                result = key + alt

        # GUI / WINDOWS
        elif cmd in ("GUI", "WINDOWS"):
            if args:
                result = (
                    DuckEncoder.keyInstr2USBBytes(
                        args, keyProp, langProp
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_LEFT_GUI",
                        keyProp,
                        langProp,
                    )
                )
            else:
                result = (
                    DuckEncoder.prop2USBByte(
                        "KEY_LEFT_GUI",
                        keyProp,
                        langProp,
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_LEFT_GUI",
                        keyProp,
                        langProp,
                    )
                )

        # COMMAND
        elif cmd == "COMMAND":
            if args:
                result = (
                    DuckEncoder.keyInstr2USBBytes(
                        args, keyProp, langProp
                    )
                    + DuckEncoder.prop2USBByte(
                        "MODIFIERKEY_LEFT_GUI",
                        keyProp,
                        langProp,
                    )
                )
            else:
                result = (
                    DuckEncoder.prop2USBByte(
                        "KEY_COMMAND",
                        keyProp,
                        langProp,
                    )
                    + b"\x00"
                )

        else:
            # Everything else is treated as a direct key input.
            result = (
                DuckEncoder.keyInstr2USBBytes(
                    cmd, keyProp, langProp
                )
                + b"\x00"
            )

        return result

    @staticmethod
    def prop2USBByte(prop, keyProp, langProp):
        keyval = None

        if prop in keyProp:
            keyval = keyProp[prop]
        elif prop in langProp:
            keyval = langProp[prop]

        if keyval is None:
            print(
                f"Error: No keycode entry for {prop}",
                file=sys.stderr,
            )
            print(
                "Warning: this could corrupt generated output file",
                file=sys.stderr,
            )
            return b""

        if keyval[0:2].upper() == "0X":
            keyval = int(keyval, 16)
        else:
            keyval = int(keyval)

        return bytes([keyval])

    @staticmethod
    def delay2USBBytes(delay):
        result = b""

        # Python 2 used integer division for "/".
        count = delay // 255
        remain = delay % 255

        for _ in range(count):
            result += b"\x00\xff"

        result += bytes([0, remain])

        return result

    @staticmethod
    def keyInstr2USBBytes(keyinstr, keyProp, langProp):
        keyval = None
        key_entry = ""

        #
        # Language fix:
        #
        # A one-character key instruction must go through ASCII
        # translation so language layouts are respected.
        #
        if len(keyinstr) == 1:
            keyval = DuckEncoder.ASCIIChar2USBBytes(
                keyinstr,
                keyProp,
                langProp,
            )

            if not keyval:
                return b""

            # ASCIIChar2USBBytes normally returns key + modifier.
            # For key instructions we need only the key byte.
            return keyval[:1]

        key_entry = "KEY_" + keyinstr.strip()

        # First attempt.
        if key_entry in keyProp:
            keyval = keyProp[key_entry]
        elif key_entry in langProp:
            keyval = langProp[key_entry]

        # Try aliases.
        if keyval is None:
            keyinstr = keyinstr.strip().upper()

            keyinstr = {
                "ESCAPE": "ESC",
                "RETURN": "ENTER",
                "DEL": "DELETE",
                "BREAK": "PAUSE",
                "CONTROL": "CTRL",
                "DOWNARROW": "DOWN",
                "UPARROW": "UP",
                "LEFTARROW": "LEFT",
                "RIGHTARROW": "RIGHT",
                "MENU": "APP",
                "WINDOWS": "GUI",
                "PLAY": "MEDIA_PLAY_PAUSE",
                "PAUSE": "MEDIA_PLAY_PAUSE",
                "STOP": "MEDIA_STOP",
                "MUTE": "MEDIA_MUTE",
                "VOLUMEUP": "MEDIA_VOLUME_INC",
                "VOLUMEDOWN": "MEDIA_VOLUME_DEC",
                "SCROLLLOCK": "SCROLL_LOCK",
                "NUMLOCK": "NUM_LOCK",
                "CAPSLOCK": "CAPS_LOCK",
            }.get(keyinstr)

            if keyinstr:
                key_entry = "KEY_" + keyinstr.strip()

                if key_entry in keyProp:
                    keyval = keyProp[key_entry]
                elif key_entry in langProp:
                    keyval = langProp[key_entry]

        if keyval is None:
            sys.stderr.write(
                f"Error: No keycode entry for {key_entry}\n"
            )
            sys.stderr.write(
                "Warning: this could corrupt generated output file\n"
            )
            return b""

        if keyval[0:2].upper() == "0X":
            keyval = int(keyval, 16)
        else:
            keyval = int(keyval)

        return bytes([keyval])

    @staticmethod
    def ASCIIChar2USBBytes(char, keyProp, langProp):
        result = b""

        val = ord(char)

        hexval = f"{val:X}"
        if len(hexval) == 1:
            hexval = "0" + hexval

        if val < 0x80:
            name = "ASCII_" + hexval
        else:
            name = "ISO_8859_1_" + hexval

        if name not in langProp:
            print(
                f"{char} interpreted as {name}, "
                "but not found in chosen language property file. "
                "Skipping character!",
                file=sys.stderr,
            )
        else:
            for key_entry in langProp[name].split(","):
                key_entry = key_entry.strip()

                keyval = None

                if key_entry in keyProp:
                    keyval = keyProp[key_entry]
                elif key_entry in langProp:
                    keyval = langProp[key_entry]

                if keyval is None:
                    print(
                        f"Error: No keycode entry for {key_entry}",
                        file=sys.stderr,
                    )
                    print(
                        "Warning: this could corrupt generated output file",
                        file=sys.stderr,
                    )
                    return b""

                if keyval[0:2].upper() == "0X":
                    keyval = int(keyval, 16)
                else:
                    keyval = int(keyval)

                result += bytes([keyval])

            # Add modifier byte if only the key byte was generated.
            if len(result) == 1:
                result += b"\x00"

        return result

    @staticmethod
    def parseScript(source, keyProp, langProp):
        result = b""

        lines = source.splitlines()

        lastLine = None

        for line in lines:
            line = line.strip()

            # Skip blank lines and comments.
            if (
                not line
                or line.startswith("//")
                or line.startswith("REM ")
            ):
                continue

            # REPEAT instruction.
            if line.startswith("REPEAT "):
                instr = line.split(" ", 1)

                if len(instr) == 1 or lastLine is None:
                    continue

                repeat_count = int(instr[1].strip())

                for _ in range(repeat_count):
                    result += DuckEncoder.parseScriptLine(
                        lastLine,
                        keyProp,
                        langProp,
                    )
            else:
                result += DuckEncoder.parseScriptLine(
                    line,
                    keyProp,
                    langProp,
                )

            lastLine = line

        return result

    @staticmethod
    def pwd():
        return os.path.dirname(os.path.abspath(__file__)) or "."

    @staticmethod
    def generatePayload(source, lang):
        script_dir = DuckEncoder.pwd()

        keyboard = DuckEncoder.readResource(
            os.path.join(
                script_dir,
                "resources",
                "keyboard.properties",
            )
        )

        language = DuckEncoder.readResource(
            os.path.join(
                script_dir,
                "resources",
                f"{lang}.properties",
            )
        )

        return DuckEncoder.parseScript(
            source,
            keyboard,
            language,
        )

    def out2hid(self, data):
        with open(self.__key_dev_file, "wb") as f:
            for i in range(0, len(data), 2):
                if i + 1 >= len(data):
                    break

                key = data[i]
                mod = data[i + 1]

                # Delay.
                if key == 0:
                    time.sleep(mod / 1000.0)
                    continue

                out = bytes(
                    [mod, 0, key, 0, 0, 0, 0, 0]
                    + [0] * 8
                )

                f.write(out)
                f.flush()

    def outhidString(self, string):
        payload = b""

        for c in string:
            payload += DuckEncoder.ASCIIChar2USBBytes(
                c,
                self.keyboard,
                self.language,
            )

        self.out2hid(payload)

    def outhidStringDirect(self, string):
        with open(self.__key_dev_file, "wb") as f:
            for c in string:
                data = DuckEncoder.ASCIIChar2USBBytes(
                    c,
                    self.keyboard,
                    self.language,
                )

                for i in range(0, len(data), 2):
                    if i + 1 >= len(data):
                        break

                    key = data[i]
                    mod = data[i + 1]

                    if key == 0:
                        time.sleep(mod / 1000.0)
                        continue

                    out = bytes(
                        [mod, 0, key, 0, 0, 0, 0, 0]
                        + [0] * 8
                    )

                    f.write(out)

                f.flush()

    def outhidDuckyScript(self, source):
        payload = DuckEncoder.parseScript(
            source,
            self.keyboard,
            self.language,
        )

        self.out2hid(payload)

    def setLanguage(self, str_lang):
        res = ""

        if self.__str_lang != str_lang:
            try:
                self.language = DuckEncoder.readResource(
                    os.path.join(
                        DuckEncoder.pwd(),
                        "resources",
                        f"{str_lang}.properties",
                    )
                )
            except OSError:
                res = (
                    f"No language file for '{str_lang}', "
                    "resetting to 'us'"
                )
                self.print_debug(res)

                self.language = DuckEncoder.readResource(
                    os.path.join(
                        DuckEncoder.pwd(),
                        "resources",
                        "us.properties",
                    )
                )

                return res

            self.__str_lang = str_lang
            res = f"language set to '{str_lang}'"
            self.print_debug(res)

        return res

    def getLanguage(self):
        return self.__str_lang

    def setKeyDevFile(self, key_dev_file):
        self.__key_dev_file = key_dev_file

    def print_debug(self, message):
        if self.DEBUG:
            print(message)

    def __init__(
        self,
        lang="us",
        key_dev_file="/dev/hidg0",
    ):
        self.DEBUG = False

        self.keyboard = DuckEncoder.readResource(
            os.path.join(
                DuckEncoder.pwd(),
                "resources",
                "keyboard.properties",
            )
        )

        self.__key_dev_file = key_dev_file
        self.__str_lang = ""

        self.setLanguage(lang)


def usage():
    usagescr = '''Duckencoder python port 1.0 by MaMe82
================================================

Creds to:
    hak5Darren for original duckencoder
    https://github.com/hak5darren/USB-Rubber-Ducky

Converts DuckyScript source into an encoded binary payload.

Usage:
    python3 duckencoder.py -i [file]
        Encode DuckyScript source given by -i file

    python3 duckencoder.py -i [file] -o [outfile]
        Encode DuckyScript source to output file

Arguments:
   -i [file]         Input file in DuckyScript format
   -o [file]         Output file for encoded payload,
                     defaults to inject.bin
   -l <layout name>  Keyboard Layout (us/fr/pt/de ...)
   -p, --passthru    Read script from stdin and print result on stdout
                     (ignore -i, -o)
   -r, --rawpassthru Read input from stdin as STRING instead of DuckyScript
   -h                Print this help screen
'''
    print(usagescr)


def main(argv):
    script_dir = (
        os.path.dirname(os.path.abspath(__file__))
        or "."
    )

    ifile = ""
    source = None
    ofile = "inject.bin"
    lang = "us"
    rawpassthru = False

    try:
        opts, args = getopt.getopt(
            argv,
            "hi:o:l:pr",
            [
                "help",
                "input=",
                "output=",
                "language=",
                "passthru",
                "rawpassthru",
            ],
        )
    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)

            # DuckyScript source is text.
            with open(
                ifile,
                "r",
                encoding="utf-8",
            ) as f:
                source = f.read()

        elif opt in ("-l", "--language"):
            lfile = os.path.join(
                script_dir,
                "resources",
                f"{arg}.properties",
            )

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

            lang = arg

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

        elif opt in ("-p", "--passthru"):
            # Read input from stdin, no output file.
            ofile = None
            source = sys.stdin.read()

        elif opt in ("-r", "--rawpassthru"):
            # Read input from stdin as raw text.
            rawpassthru = True
            ofile = None
            source = sys.stdin.read()

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

    if rawpassthru:
        result = b""

        keyboard = DuckEncoder.readResource(
            os.path.join(
                script_dir,
                "resources",
                "keyboard.properties",
            )
        )

        language = DuckEncoder.readResource(
            os.path.join(
                script_dir,
                "resources",
                f"{lang}.properties",
            )
        )

        for line in source.splitlines():
            for c in line:
                keydata = DuckEncoder.ASCIIChar2USBBytes(
                    c,
                    keyboard,
                    language,
                )

                if keydata:
                    result += keydata

    else:
        result = DuckEncoder.generatePayload(
            source,
            lang,
        )

    if ofile is None:
        # Binary payload to stdout.
        sys.stdout.buffer.write(result)
    else:
        with open(ofile, "wb") as f:
            f.write(result)


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

    main(sys.argv[1:])