symkey/capture.py

104 lines
2.2 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)
now = time.time()
_do_delay(args.delay)
_do_trigger(args.trigger)
# Add dummy events to lock in the time start
event_queue = queue.Queue()
event_queue.put(
mouse.ButtonEvent(
event_type=mouse.UP,
button=mouse.RIGHT,
time=now,
))
event_queue.put(
keyboard.KeyboardEvent(
event_type=keyboard.KEY_UP,
name=" ",
scan_code=57,
time=now,
)
)
hook = functools.partial(_on_hook, event_queue)
keyhook = keyboard.hook(hook)
mousehook = mouse.hook(hook)
print("Capturing...")
try:
keyboard.wait(args.trigger)
except KeyboardInterrupt:
keyboard.unhook(keyhook)
mouse.unhook(mousehook)
_save_events(event_queue, args.output)
def _on_hook(event_queue, event):
LOGGER.debug(str(event))
event_queue.put(event, block=False)
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()