104 lines
2.2 KiB
Python
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()
|