My solution to Inaction from SuSeC CTF 2020

solve_inaction.py

"""inaction_solve.py

CTF: SuSeC CTF 2020 - https://ctftime.org/event/1007
Challenge: Inaction - https://ctftime.org/task/10729

Flag: SUSEC{5m4Rt___p0W___cH4lL3n93}

twitter.com/0xTowel
"""
import re

from sys import exit, stderr
from typing import List
from pwn import remote


####################
# Global Constants #
####################

# Connection Information
SERVER = '66.172.27.57'
PORT = 3114

# Code constants
RE_INT = re.compile(br'\d+')
P = 2**2281 - 1
PAD = 2228


def solve(val: int, i: int) -> str:
    """Solve the challenge for i iterations.

    Checks both candidate roots for the next set of roots.
    If one deadends, we use the other.
    """
    for _ in range(i):
        vals = prime_mod_sqrt(val, P)
        candidate_1, candidate_2 = min(vals) ^ PAD, max(vals) ^ PAD
        check = prime_mod_sqrt(candidate_1, P)
        val = candidate_1 if check != [] else candidate_2
    return str(min(val, (val * -1) % P))


def main() -> None:
    """Main challenge interaction.

    We connect, and try and solve 40 challenges with 40 iterations
    each.
    """
    r = remote(SERVER, PORT)
    r.recvuntil(b'-\ny')  # Server welcome message

    # Solve 40 challenges with 40 iterations each
    stage = 40
    for _ in range(stage):
        y = grab_int(r.recvline())
        sol = solve(y, stage)
        r.sendlineafter(b'solution:', str(sol))
        print(r.recvlines(numlines=2)[1].decode())
    r.close()


#####################
# Utility Functions #
#####################

def grab_int(data: bytes) -> int:
    """Grab an integer from a bytestring."""
    search_results = RE_INT.search(data)
    if search_results:
        return int(search_results.group())
    print(f'[!] Error parsing server challenge: {data}', file=stderr)
    exit(-1)

# Sourced from stack overflow because I don't wanna
# reinvent the wheel during a CTF :P
#
# https://codereview.stackexchange.com/questions/
# 43210/tonelli-shanks-algorithm-implementation-of-prime-modular-square-root

def legendre_symbol(a: int, p: int) -> int:
    """Legendre symbol

    Defined if a is a quadratic residue modulo odd prime.
    """
    ls = pow(a, (p - 1) // 2, p)
    if ls == p - 1:
        return -1
    return ls


def prime_mod_sqrt(a: int, p: int) -> List[int]:
    """Square root modulo prime number

    Solve the equation: x^2 = a mod p
    and return list of x solution
    """
    a %= p

    # Simple case
    if a == 0:
        return [0]
    if p == 2:
        return [a]

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

    # Simple case
    if p % 4 == 3:
        x = pow(a, (p + 1) // 4, p)
        return [x, p - x]

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

    # 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
        i, e = 0, 2
        for i in range(1, m):
            if pow(t, e, p) == 1:
                break
            e *= 2

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

    return [x, p - x]


if __name__ == '__main__':
    main()