import socket
import json
import time
import random
import string
import argparse

KEYS_TO_USE = string.ascii_lowercase + string.digits

def qmp_execute(sock, command, args=None):
    payload = {"execute": command}
    if args:
        payload["arguments"] = args
    sock.sendall(json.dumps(payload).encode())
    return json.loads(sock.recv(1024))

def main(qmp_socket, interval):
    try:
        # connect
        client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client.connect(qmp_socket)
        
        # qmp handshake
        client.recv(1024) 
        qmp_execute(client, "qmp_capabilities")
        
        print(f"Connected. Spamming random keys every {interval}s...")

        while True:
            key = random.choice(KEYS_TO_USE)
            
            qmp_execute(client, "send-key", {"keys": [{"type": "qcode", "data": key}]})
            
            print(f"Sent: {key}", end="\r")
            time.sleep(interval)

    except KeyboardInterrupt:
        print("\nStopping stress test.")
    except Exception as e:
        print(f"\nError: {e}")
    finally:
        client.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--qmp-socket", help="Path to QMP socket file", type=str, default="/tmp/qmp-sock")
    parser.add_argument("--interval", help="PS/2 keyboard torture interval", type=float, default=0.05)

    args = parser.parse_args()

    main(args.qmp_socket, args.interval)