#!/usr/bin/python2
from pwn import process, remote
import gmpy
import lcgcrackmultiples
import time

# Linear Congruential Generator
def lcg(seed, a, c, m):
    last = seed
    while True:
        yield last
        last = (a * last + c) % m

def get_next(prng):
    prng.recv(1024)
    prng.sendline('2')
    return int(prng.recvline())

def send_next(prng, n):
    prng.sendline('1')
    prng.recv(1024)
    prng.sendline(str(n))
    print(prng.recv(1024))

def get_seq(prng, sample):
    seq = [get_next(prng) for k in range(sample)]
    print('Got this sequence to analyze: ')
    print(seq)
    return seq

def get_prng(opt, data):
    if opt == 1:
        return process(data['path'])
    else:
        return remote(data['host'], data['port'])

def my_prng():
    return get_prng(1, {'path': 'lcg/lcg'})

def remote_prng():
    return get_prng(2, {
        'host': '172.16.58.3',
        'port': 7412
    })

# Equation from 3 sequence values:
# 842389455, 3301052331, 1833279318

# 3301052331 = a*842389455 + c (mod m)
# 1833279318 = a*3301052331 + c (mod m)
# 3301052331 - 1833279318 = a * (842389455 - 3301052331) (mod m)

# a * -2458662876 = 1467773013 (mod m)
# This equation, solved for a, can be expressed as:
def get_a(m, seq):
    k = (seq[0] - seq[1]) % m
    inv = int(gmpy.invert(k, m))

    return (inv * (seq[1] - seq[2])) % m

def check_values(prng,a,c,n,seq):
    # Create an LCG with initial value = seq[0]
    # And test whether it produces the same values
    candidate_lcg = lcg(seq[0], a, c, n)
    flag = True
    for i in seq:
        if next(candidate_lcg) != i:
            flag = False
    # If so, show values for N, A and C
    # Then generate ten next values and send
    if flag == True:
        print('[M,A,C]', [n, a, c])
        for k in range(10):
            num = next(candidate_lcg)
            print('[SEND]', num)
            send_next(prng, num)

        print(prng.recvall())
        return True
    return False

# Test a bunch of modulus
def find_modulus(prng, seq, start, n_tests):
    for n in range(start, start+n_tests):
        a = get_a(n, seq)
        if a != 0: # If an a exists
            # Figure out the c for this a
            c = (seq[1] - seq[0]*a) % n
            if (check_values(prng,a,c,n,seq) == True):
                return True
    return False

def crackwithmultiples(prng,seq):
    try:
        modulus, multiplier, increment = lcgcrackmultiples.crack(seq)
        return check_values(prng,multiplier,increment,modulus,seq)
    except:
        return False

# Observed maximum value after leaking a bunch of numbers
# (modulus must be close!)
max_seen = 4294915818
# Amount of numbers to test after the observed maximum
n_tests = 1000000
# Sequence from remote PRNG
# Sample should be any number > 3
# (3 to solve the equation, other values for validation)
sample_n = 6

# Setting up PRNG
prng = my_prng()
# Sequence from remote PRNG
seq = get_seq(prng, sample_n)

start = time.time()
print("\nMethod using multiples (cf. https://tailcall.net/blog/cracking-randomness-lcgs/):")
done = crackwithmultiples(prng,seq)
while not done:
    print('[NOT FOUND] Modular inverse not found!')
    print('Fetching another sequence and trying again.')
    try:
        prng.close()
    except IOError:
        print("Problem closing prng")
    prng = my_prng()
    seq = get_seq(prng, sample_n)
    done = crackwithmultiples(prng,seq)
end = time.time()
print("Took: ",end-start)
prng.close()

prng = my_prng()
seq = get_seq(prng, sample_n)

start = time.time()
print("\nMethod using bruteforce on modulus:")
done = find_modulus(prng, seq, max_seen, n_tests)
while not done:
    print('[NOT FOUND] Bad seed or modulus not in range!')
    print('Fetching another sequence and trying again.')
    try:
        prng.close()
    except IOError:
        print("Problem closing prng")
    prng = my_prng()
    seq = get_seq(prng, sample_n)
    done = find_modulus(prng, seq, max_seen, n_tests)
end = time.time()
print("Took: ",end-start)
prng.close()

