#!/usr/bin/env python
import os
import sys
import time
import json
from subprocess import call

# This python script is heavily based on DuckToolkit (https://ducktoolkit.com),
# thanks to DuckToolkit developers Kevin Breen and James Hall.

__description__ = 'Interpreter of Ducky Scripts'
__author__ = 'Bash Bunny Development Team'
__version__ = '0.4'
__date__ = '03/27/2017'

LANG_DIR = "/usr/local/bunny/lib/languages"

def hidg_write(elements):
    values = bytearray(elements)
    not_hold = bytearray([0, 0, 0, 0, 0, 0, 0, 0])

    hidg = open("/dev/hidg0", "wb")
    hidg.write(values)
    hidg.write(not_hold)
    hidg.close()


def run_ducky_line(context, line, lang_file):
    # REM Comments
    if line.startswith('REM') or line.startswith('rem'):
        return context

    # Last Command
    last_command = context[1]
    last_instruction = context[2]
    default_delay = context[3]

    line = line.lstrip().rstrip()
    parsed_line = line.split(' ', 1)

    if len(parsed_line) >= 2:
        cmd = parsed_line[0].strip()
        instruction = parsed_line[1].rstrip()
    else:
        cmd = parsed_line[0].strip()
        instruction = False

    # Default Delay
    if cmd in ['DEFAULT_DELAY', 'DEFAULTDELAY']:
        try:
            default_delay = int(instruction)
            context_ret = ((1, ''), cmd, instruction, default_delay)
            return context_ret
        except Exception as e:
            error_msg = 'Handling default delay error'
            context_ret = ((0, error_msg), '', '', 0)
            return context_ret

    # Repeat
    repeat_count = 1
    if cmd in ['REPEAT', 'repeat', 'REPLAY', 'replay']:
        try:
            repeat_count = int(instruction)
        except Exception as e:
            print e
            error_line = 'Repeat value not valid'

        cmd = last_command
        instruction = last_instruction

    for i in range(repeat_count):
        if cmd == 'KEYCODE':
            for keycode in instruction.split(" "):
                elements = keycode.split(",")
                elements = [int(i, 16) for i in elements]
                for i in range(5):
                    elements.append(0)
                hidg_write(elements)

        elif cmd == 'STRING':
            for char in instruction:
                elements = lang_file[char].split(",");
                elements = [int(i, 16) for i in elements]
                for i in range(5):
                    elements.append(0)
                hidg_write(elements)

        elif cmd == 'DELAY':
            time.sleep(0.001 * int(instruction));

        elif cmd in lang_file.iterkeys():
            elements = lang_file[cmd].split(",");
            elements = [int(i, 16) for i in elements]
            for i in range(5):
                elements.append(0)

            if instruction:
                param = lang_file[instruction].split(",");
                param = [int(i, 16) for i in param]
                elements[0] |= param[0]
                elements[2] |= param[2]

            hidg_write(elements)

        else:
            err_line = "Command {0} Not in Language File".format(cmd)
            print "err_line = ", err_line
            context_ret = ((0, err_line), '', '', 0)
            return context_ret

        # Add Default Delay
        if default_delay:
            time.sleep(0.001 * int(default_delay))

    # Return normally
    context_ret = ((1, ''), cmd, instruction, default_delay)
    return context_ret


def run_ducky_payload(duck_filename, lang_file):
#   print "duck_filename = ", duck_filename
    try:
        duck_text = open(duck_filename, 'rb').read()
    except Exception as e:
        print "  [!] Unable to open input file: {0}".format(duck_filename)
        sys.exit()

    duck_text = duck_text.replace("\r", "")
    context = ((0, ''), '', '', 0)

    for line in duck_text.split('\n'):
        if len(line) > 0:
            context = run_ducky_line(context, line, lang_file)


def run_script(input_line, duck_lang):
    language_dict = os.path.join(LANG_DIR, '{0}.json'.format(duck_lang))
    lang_file = json.load(open(language_dict))

    # define a tuple:
    # 0: return tuple [value (int), message (string)]
    # 1: last_cmd (string)
    # 2: last_instruction (string)
    # 3: default_delay (int)
    context = ((0, ''), '', '', 0)

    fullpath = "/root/udisk/payloads/"
    input_line = input_line.replace("\r", "")
    line = input_line
    if len(line) > 0:
        if line.rstrip()[-4:] == '.txt':
            filename = line.rstrip()
            run_ducky_payload(fullpath+filename, lang_file)
            # clear context after 'Q payload.txt'
            context = ((0, ''), '', '', 0)
        else:
            # run standard ducky script
            context = run_ducky_line(context, line, lang_file)


def list_languages():
    languages = []
    for filename in os.listdir(LANG_DIR):
        if filename.endswith('.json'):
            languages.append(filename)
    return languages


if __name__ == "__main__":

    language = os.getenv("DUCKY_LANG", default="us").lower()

    if "{0}.json".format(language) not in list_languages():
        print "[!] Language '{0}; is not supported at this time.".format(language)
        print "[+] Supported Languages"
        for lang in list_languages():
            print "  [-] {0}".format(lang.split('.')[0])
        sys.exit()

    input_line = " ".join(sys.argv[1:])

    run_script(input_line, language)
