uberducky.c

/*
 * Copyright 2019 Mike Ryan
 *
 * This file is part of Uberducky and is released under the terms of the
 * GPL version 2. Refer to COPYING for more information.
 */

#include "hid.h"
#include "ble.h"

#include "type.h"
#include "ubertooth.h"
#include "usbapi.h"

#include <string.h>

// magic string that must be present in trigger packets
// derived from random UUID:
// fd123ff9-9e30-45b2-af0d-b85b7d2dc80c
uint8_t ble_magic[16] = {
    0x0c, 0xc8, 0x2d, 0x7d, 0x5b, 0xb8, 0x0d, 0xaf,
    0xb2, 0x45, 0x30, 0x9e, 0xf9, 0x3f, 0x12, 0xfd,
};

// magic string that triggers bootloader mode
// random UUID:
// 344bc7f2-5619-4953-9be8-9888fe29d996
uint8_t bootloader_magic[16] = {
    0x96, 0xd9, 0x29, 0xfe, 0x88, 0x98, 0xe8, 0x9b,
    0x53, 0x49, 0x19, 0x56, 0xf2, 0xc7, 0x4b, 0x34,
};

extern uint8_t script[]; // auto-generated from duckyscript input

// times in ms
#define LED_PERIOD      600
#define LED_ON_TIME     100

#define INTR_IN_EP      0x81

#define LE_WORD(x)      ((x)&0xFF),((x)>>8)
#define READ_LE(x)      ((script[x+1] << 8) | script[x])

#define NOW T0TC

// script state
#define ST_IDLE     0
#define ST_READY    1
#define ST_RUNNING  2

// run state
#define R_IDLE      0
#define R_KEY_DOWN  1
#define R_DELAY     2
#define R_STRING    3

// string state
#define S_IDLE      0
#define S_KEY_DOWN  1

int script_state = ST_IDLE;
int run_state = R_IDLE;
int string_state = S_IDLE;

unsigned script_len = 0;
unsigned script_pos = 0;
unsigned string_len = 0;
unsigned string_pos = 0;

uint32_t repeat_counter = 0;
unsigned repeat_pos = 0;
int repeating = 0;

// opcodes
//
// encoding: <op> [<arg> .. ]
// 
// NOP - no args
// KEY - key type, modifier, character (each 1 byte)
// DELAY - delay in ms (16 bit little endian)
// STRING - length (16 bit little endian), chars
// REPEAT - repeat previous command
#define OP_NOP    0
#define OP_KEY    1
#define OP_DELAY  2
#define OP_STRING 3
#define OP_REPEAT 4

#define DELAY(X) OP_DELAY, LE_WORD(X)

// demo script - print hello world
/* this is now loaded from an autogenerated .c file
uint8_t script[] = {
    LE_WORD(24),  // script len
    OP_STRING,    // string
    LE_WORD(6),   // string len
    'h', 'e', 'l', 'l', 'o', ' ',
    DELAY(1000),  // delay
    OP_STRING,    // string
    LE_WORD(5),
    'w', 'o', 'r', 'l', 'd',
    OP_KEY,       // key - enter with no modifier
    K_ENTER, 0x00, 0x00,
};
*/

static void timer0_start(void) {
    T0TCR = TCR_Counter_Reset;
    T0PR = 50000 - 1; // 1 ms
    T0TCR = TCR_Counter_Enable;

    // set up interrupt handler
    ISER0 = ISER0_ISE_TIMER0;
}

static void timer0_set_match(uint32_t match) {
    T0MR0 = match;
    T0MCR |= TMCR_MR0I;
}

static void timer0_clear_match(void) {
    T0MCR &= ~TMCR_MR0I;
}

void TIMER0_IRQHandler(void) {
    uint8_t report[8] = { 0, };
    keystroke_t next_key = { 0, };

    if (T0IR & TIR_MR0_Interrupt) {
        // ack the interrupt
        T0IR = TIR_MR0_Interrupt;

        if (script_state == ST_READY) {
            script_pos = 0;
            script_state = ST_RUNNING;
            run_state = R_IDLE;

            script_len = READ_LE(script_pos);
            script_pos += 2;

            // re-enter in 1 ms
            timer0_set_match(NOW + 1);
        } else if (script_state == ST_RUNNING) {
            uint8_t opcode;
            uint16_t delay;
            switch (run_state) {
                // idle -- get next opcode
                case R_IDLE:
                    if (script_pos >= script_len) {
                        script_state = ST_IDLE;
                        return;
                    }

                    if (repeating) {
                        if (repeat_counter == 0) {
                            repeating = 0;
                            script_pos += 3; // skip repeat opcode
                        } else {
                            --repeat_counter;
                            script_pos = repeat_pos; // jump back to prev op
                        }
                    }

                    opcode = script[script_pos++];
                    if (opcode != OP_REPEAT)
                        repeat_pos = script_pos-1;

                    switch (opcode) {
                        case OP_NOP:
                            ++script_pos;
                            timer0_set_match(NOW + 1); // re-enter in 1 ms
                            return;

                        case OP_KEY:
                            // TODO bounds check
                            next_key.type = script[script_pos++];
                            next_key.mod = script[script_pos++];
                            next_key.chr = script[script_pos++];

                            // encode and inject key
                            hid_encode(&next_key, report);
                            USBHwEPWrite(INTR_IN_EP, report, 8);
                            run_state = R_KEY_DOWN;
                            timer0_set_match(NOW + DOWN_TIME);
                            return;

                        case OP_DELAY:
                            // TODO bounds check
                            delay = READ_LE(script_pos);
                            script_pos += 2;
                            run_state = R_DELAY;
                            timer0_set_match(NOW + delay);
                            return;

                        case OP_STRING:
                            // TODO bounds check
                            string_len = READ_LE(script_pos);
                            string_pos = 0;
                            script_pos += 2;
                            run_state = R_STRING;
                            string_state = S_IDLE;
                            timer0_set_match(NOW + 1); // re-enter in 1 ms
                            return;

                        case OP_REPEAT:
                            repeating = 1;
                            repeat_counter = READ_LE(script_pos);
                            timer0_set_match(NOW + 1); // re-enter in 1 ms
                            return;
                    }
                    return;

                // key down - lift key
                case R_KEY_DOWN:
                    // all keys up
                    USBHwEPWrite(INTR_IN_EP, report, 8);
                    timer0_set_match(NOW + DOWN_TIME);
                    run_state = R_IDLE;
                    // timer0_set_match(NOW + 1);
                    return;

                case R_DELAY:
                    run_state = R_IDLE;
                    timer0_set_match(NOW + 1);
                    return;

                // string - sub state machine
                case R_STRING:
                    switch (string_state) {
                        // idle - get next key
                        case S_IDLE:
                            // end of string, next
                            if (string_pos >= string_len) {
                                script_pos += string_len;
                                run_state = R_IDLE;
                                timer0_set_match(NOW + 1);
                                return;
                            }

                            next_key.type = K_CHAR;
                            next_key.mod = 0;
                            next_key.chr = script[script_pos + string_pos];
                            ++string_pos;

                            hid_encode(&next_key, report);
                            USBHwEPWrite(INTR_IN_EP, report, 8);
                            string_state = S_KEY_DOWN;
                            timer0_set_match(NOW + DOWN_TIME);
                            return;

                        // key down - lift and go back to idle state
                        case S_KEY_DOWN:
                            USBHwEPWrite(INTR_IN_EP, report, 8);
                            timer0_set_match(NOW + DOWN_TIME);
                            string_state = S_IDLE;
                            return;
                    }
                    return;
            }
        }
    }

    // LEDs
    if (T0IR & TIR_MR1_Interrupt) {
        T0IR = TIR_MR1_Interrupt;
        RXLED_CLR;
    }
}

int magic_present(uint8_t *packet, uint8_t *magic) {
    unsigned i;

    for (i = 0; i < BLE_PACKET_SIZE - 16; ++i)
        if (memcmp(packet + i, magic, 16) == 0)
            return 1;

    return 0;
}

// from usb.c
void usb_init(void);

int main() {
    uint8_t ble_packet[BLE_PACKET_SIZE];
    int led_state = 0;
    uint32_t led_next_event = LED_PERIOD - LED_ON_TIME;

    ubertooth_init();

    timer0_start();

    usb_init();
    ble_init();

    // call USB interrupt handler continuously
    while (1) {
        USBHwISR();

        // fetch BLE packets
        if (ble_get_packet(ble_packet)) {
            // blink LED - TODO something more interesting
            RXLED_SET;
            T0MR1 = NOW + 10;
            T0MCR |= TMCR_MR1I;

            // launch script if magic string is in packet and we're idle
            if (script_state == ST_IDLE && magic_present(ble_packet, ble_magic)) {
                script_state = ST_READY;
                timer0_set_match(NOW + 1);
            }

            // if the bootloader magic is present, reset to bootloader
            else if (magic_present(ble_packet, bootloader_magic)) {
                // turn off radio
                cc2400_strobe(SRFOFF);
                while ((cc2400_status() & FS_LOCK)); // need to wait for unlock?
#ifdef UBERTOOTH_ONE
                PAEN_CLR;
                HGM_CLR;
#endif

                // reset
                bootloader_ctrl = DFU_MODE;
                reset();
            }
        }

        // blink LED
        if (NOW >= led_next_event) {
            if (led_state == 0) {
                led_state = 1;
                TXLED_SET;
                led_next_event += LED_ON_TIME;
            } else {
                led_state = 0;
                TXLED_CLR;
                led_next_event += LED_PERIOD - LED_ON_TIME;
            }
        }
    }

    return 0;
}