symkey/capture.py

import argparse
import functools
import keyboard
import logging
import mouse
import pickle
import queue
import sys
import threading
import time

LOGGER = logging.getLogger("capture")

def main() -> None:
	parser = argparse.ArgumentParser()
	parser.add_argument(
		"-d", "--delay",
		default=0,
		type=int,
		help="Seconds to wait before capture."
	)
	parser.add_argument(
		"-t", "--trigger",
		default=None,
		help="The key to use to trigger start and stop of capture."
	)
	parser.add_argument(
		"-o", "--output",
		default="dump.capture",
		help="Name of the file to capture to."
	)
	parser.add_argument(
		"--verbose",
		action="store_true",
		help="Show verbose logging.",
	)
	args = parser.parse_args()
	logging.basicConfig(
		level=logging.DEBUG if args.verbose else logging.INFO)
	if args.delay and args.trigger:
		print("You cannot specify 'delay' and 'trigger'")
		sys.exit(1)

	_do_delay(args.delay)
	_do_trigger(args.trigger)
	LOGGER.info("Capturing...")
	now = time.time()
	# Add dummy events to lock in the time start
	with open(args.output, "w") as output:
		hook = functools.partial(_on_hook, now, output)
		keyhook = keyboard.hook(hook)
		mousehook = mouse.hook(hook)
		try:
			keyboard.wait(args.trigger)
		except KeyboardInterrupt:
			keyboard.unhook(keyhook)
			mouse.unhook(mousehook)
	LOGGER.info("Wrote %s", args.output)

def _on_hook(start, output, event):
	LOGGER.debug(str(event))
	relative_time = event.time - start
	if isinstance(event, keyboard.KeyboardEvent):
		output.write(
			f"{relative_time},k,{event.event_type},{event.scan_code},{event.name}\n")
	elif isinstance(event, mouse.ButtonEvent):
		output.write(
			f"{relative_time},mb,{event.event_type},{event.button}\n")
	elif isinstance(event, mouse.MoveEvent):
		output.write(
			f"{relative_time},mm,{event.x},{event.y}\n")
	elif isinstance(event, mouse.WheelEvent):
		output.write(
			f"{relative_time},mw,{event.delta}\n")
	else:
		raise ValueError(f"{event} is not recognized")

def _do_delay(delay: int) -> None:
	if not delay:
		return
	print("\n")
	for i in range(delay):
		print(f"\rStarting in {delay-i} seconds")
		time.sleep(1)

def _do_trigger(trigger: str) -> None:
	if not trigger:
		return
	print(f"Waiting for '{trigger}'")
	keyboard.wait(trigger)

def _save_events(event_q: queue.Queue, filename: str) -> None:
	events = []
	while not event_q.empty():
		event = event_q.get(block=False)
		events.append(event)
	with open(filename, "wb") as output:
		pickle.dump(events, output)
	print(f"Wrote to {filename}")

if __name__ == "__main__":
	main()