CCCamp - Prejudiced randomness

solve.py

#!/usr/bin/env python2.7
import gmpy
import random
from pwn import *

def legendre_symbol(a, p):
    ls = pow(a, (p - 1)/2, p)
    if ls == p - 1:
        return -1
    return ls

def _prime_mod_sqrt(a, p):
    # Reduce a mod p
    a %= p

    # Special case which solution is simple
    if p % 4 == 3:
        x = pow(a, (p + 1)/4, p)
        return x

    # Factor p-1 on the form q * 2^s (with Q odd)
    q, s = p - 1, 0
    while q % 2 == 0:
        q, s = q >> 1, s + 1

    # Select a z which is a quadratic non resudue modulo p
    z = 1
    while legendre_symbol(z, p) != -1:
        z += 1
    c = pow(z, q, p)

    # Search for a solution
    x = pow(a, (q + 1)/2, p)
    t = pow(a, q, p)
    m = s
    while t != 1:
        # Find the lowest i such that t^(2^i) = 1
        e = 2
        for i in xrange(1, m):
            if pow(t, e, p) == 1:
                break
            e *= 2

        # Update next value to iterate
        b = pow(c, 1 << (m - i - 1), p)
        x = (x * b) % p
        t = (t * b * b) % p
        c = (b * b) % p
        m = i

    return x

def prime_mod_sqrt(a, p, e=1):
    # Reduce a mod p^e
    a %= p**e

    # Handle prime 2 special case
    if p == 2:
        if e >= 3 and a % 8 == 1:
            res = []
            for x in [1, 3]:
                for k in xrange(3, e):
                    i = (x*x - a)/(2**k) % 2
                    x = x + i*2**(k-1)
                res.append(x)
                res.append(p**e - x)
            return res
        # No solution if a is odd and a % 8 != 1
        if e >= 3 and a % 2 == 1:
            return []
        # Force brut if a is even or e < 3 (for now)
        return [x for x in xrange(0, p**e) if x*x % p**e == a % p**e]

    # Check solution existence on odd prime
    ls = legendre_symbol(a, p)
    if ls == -1:
        return []

    # Case were a is 0 or p multiple
    if ls == 0:
        if a % p**e == 0:
            return [0]
        return []

    # Hensel lemma lifting from x^2 = a mod p solution
    x  = _prime_mod_sqrt(a, p)
    for i in xrange(1, e):
        f = x*x - a
        fd = 2*x
        t = - (f / p**i) * gmpy.invert(fd, p**i)
        x = x + t * p**i % p**(i+1)
    return [x, p**e - x]

io = remote('hax.allesctf.net', 7331)
done = False

for i in range(42):
    p  = gmpy.next_prime(random.randint(2**512, 2**513))
    n  = p * p

    io.recvuntil('> ')
    io.sendline(str(n))

    io.recvuntil('root of\n')
    r = long(io.recvline().strip(), 10)

    # compute any root and send it
    roots   = prime_mod_sqrt(r,p,e=2)
    anyroot = int(roots[0])

    io.sendline(str(anyroot))
    io.recvuntil(': ')

    line = io.recvline()
    assert('you win' in line), 'you did not win'
    log.success('result: %s' % line)

    io.recvuntil(': ')
    io.sendline(str(p))
    io.recvuntil(': ')
    io.sendline(str(p))

io.recvline()
io.recvline()
log.success('flag: %s' % io.recvline().strip())