symkey/capture.py

103 lines
2.5 KiB
Python

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()