#!/usr/bin/env python3

BUILDDIR='@BUILDDIR@'
SRCDIR='@SRCDIR@'

import signal
import argparse
import os
import sys
library_path = BUILDDIR + '/libfprint/'

# Relaunch ourselves with a changed environment so
# that we're loading the development version of libfprint
if 'LD_LIBRARY_PATH' not in os.environ or not library_path in os.environ['LD_LIBRARY_PATH']:
    os.environ['LD_LIBRARY_PATH'] = library_path
    os.environ['GI_TYPELIB_PATH'] = f'{BUILDDIR}/libfprint/'
    os.environ['FP_DEVICE_EMULATION'] = '1'
    try:
        os.execv(sys.argv[0], sys.argv)
    except Exception as e:
        print('Could not run script with new library path')
        sys.exit(1)

import gi
gi.require_version('FPrint', '2.0')
from gi.repository import FPrint

gi.require_version('GUsb', '1.0')
from gi.repository import GUsb

import shutil
import subprocess
import tempfile
import time

def test_variant(value):
    if (not value or any(not (char.islower() or char == '-') for char in value) or
        value.startswith('-') or value.endswith('-')):
        raise argparse.ArgumentTypeError(
            'must contain only lowercase letters and dashes')
    return value

parser = argparse.ArgumentParser(
    description='Capture USB traffic for a libfprint driver test.')
parser.add_argument('driver')
parser.add_argument('variant', nargs='?', type=test_variant,
                    help='Optional, lowercase test variant name')
parser.add_argument('--test', choices=('capture', 'custom'),
                    help='Run only the selected test (default: run both)')
options = parser.parse_args()

driver_name = options.driver
os.environ['FP_DRIVERS_ALLOWLIST'] = driver_name

test_variant = options.variant

# Check that running as root

if os.geteuid() != 0:
    print(f'{sys.argv[0]} is expected to be run as root')
    sys.exit(1)

# Check that tshark is available

tshark = shutil.which('tshark')
if not tshark:
    print("The 'tshark' WireShark command-line tool must be installed to capture USB traffic")
    sys.exit(1)

# Find the fingerprint reader
ctx = FPrint.Context()
ctx.enumerate()
devices = ctx.get_devices()
if len(devices) == 0:
    print('Could not find a supported fingerprint reader')
    sys.exit(1)
elif len(devices) > 1:
    print('Capture requires a single supported fingerprint reader to be plugged in')
    sys.exit(1)

test_name = driver_name
if test_variant:
    test_name = driver_name + '-' + test_variant
usb_device = devices[0].get_property('fpi-usb-device')
bus_num = usb_device.get_bus()
device_num = usb_device.get_address()

print(f'### Detected USB device /dev/bus/usb/{bus_num:03d}/{device_num:03d}')

# Make directory

test_dir = SRCDIR + '/tests/' + test_name
os.makedirs(test_dir, mode=0o775, exist_ok=True)
custom_script = os.path.join(test_dir, 'custom.py')
if options.test == 'custom' and not os.path.exists(custom_script):
    parser.error(f'custom test script does not exist: {custom_script}')

# Capture device info

args = ['umockdev-record', f'/dev/bus/usb/{bus_num:03d}/{device_num:03d}']
device_out = open(test_dir + '/device', 'w')
process = subprocess.Popen(args, stdout=device_out)
process.wait()

# Run capture
# https://osqa-ask.wireshark.org/questions/53919/how-can-i-precisely-specify-a-usb-device-to-capture-with-tshark/

print(f'### Reseting USB port (as descriptors could be missing in the dump otherwise)')
usb_device.open()
usb_device.reset()
usb_device.close()

def capture_test(cmd, capture_file, prompt):
    print(f'### Starting USB capture on usbmon{bus_num}')
    capture_pid = os.fork()
    assert(capture_pid >= 0)

    unfiltered_cap_path = os.path.join(tempfile.gettempdir(),
                                       f'{capture_file}-unfiltered.pcapng')
    if capture_pid == 0:
        os.setpgrp()
        args = ['tshark', '-q', '-i', f'usbmon{bus_num}', '-w', unfiltered_cap_path]
        os.execv(tshark, args)

    # Wait 1 sec to settle (we can assume setpgrp happened)
    time.sleep(1)

    print(prompt)
    with subprocess.Popen([sys.executable] + cmd) as capture_process:
        capture_process.wait()
        if capture_process.returncode != 0:
            print('Failed to capture fingerprint')
            os.killpg(capture_pid, signal.SIGKILL)
            sys.exit(1)

    os.kill(capture_pid, signal.SIGTERM)
    try:
        r = t_waitpid(capture_pid, 2)
        # Kill if nothing died
        if r[0] == 0:
            os.kill(capture_pid, signal.SIGKILL)
    except ChildProcessError:
        pass

    try:
        while True:
            r = t_waitpid(-capture_pid, timeout=2)
            # Kill the process group, if nothing died (and there are children)
            if r[0] == 0:
                os.killpg(capture_pid, signal.SIGKILL)
    except ChildProcessError:
        pass

    # Filter the capture
    print(f'\n### Saving USB capture as test case {test_name}')
    args = ['tshark', '-r', unfiltered_cap_path,
            '-Y', f'usb.bus_id == {bus_num} and usb.device_address == {device_num}',
            '-w', os.path.join(test_dir, capture_file)]
    with subprocess.Popen(args, stderr=subprocess.DEVNULL) as filter_process:
        filter_process.wait()

def t_waitpid(pid, timeout):
    timeout = time.time() + timeout
    r = os.waitpid(pid, os.WNOHANG)
    while timeout > time.time() and r[0] == 0:
        time.sleep(0.1)
        r = os.waitpid(pid, os.WNOHANG)

    return r

if options.test in (None, 'capture'):
    capture_test([SRCDIR + '/tests/capture.py', test_dir + '/capture.png'],
                 'capture.pcapng',
                 '### Capturing fingerprint, please swipe or press your finger on the reader')

if options.test in (None, 'custom') and os.path.exists(custom_script):
    capture_test([custom_script], 'custom.pcapng',
                 '### Running the custom fingerprint capture')

print(f"\nDone! Don't forget to add {test_name} to tests/meson.build")