"""
EXAMPLE: ASIS Finals 2015 — "Bodu" (Crypto, 175p)
https://github.com/p4-team/ctf/tree/master/2015-10-10-asisfin/crypto_175_bodu

VULN: RSA with a VERY large public exponent e and small private exponent d
      (Boneh-Durfee / Coppersmith small-d attack, d < n^0.292). Wiener fails.
      Faithful port of jvdsn/crypto-attacks boneh_durfee (Herrmann-May bivariate
      small-roots) to pure Python. The bivariate polynomial lives in the
      quotient ring u = 1 + x*y, so every monomial is kept in reduced form
      (no term has both x and y non-zero). Lattice reduced with fpylll LLL;
      roots found with sympy resultants.

Run:
    cd /home/code/ctfkit
    python3 examples/bodu/solve.py
Expected flag: ASIS{b472266d4dd916a23a7b0deb5bc5e63f}
"""
import sys
import math
from math import comb
from Crypto.PublicKey import RSA
from Crypto.Util.number import long_to_bytes, isPrime

sys.path.insert(0, "/home/code/ctfkit")
from fpylll import IntegerMatrix, LLL
import sympy as sp

HERE = __file__.rsplit("/", 1)[0]

# Reduced monomial basis: (au, ax, ay) with ax==0 or ay==0  (x*y -> u-1).


def _reduce_one(au, ax, ay):
    """Reduce x^ax y^ay u^au using x*y = u-1. Returns dict of reduced monomials."""
    if ax == 0 or ay == 0:
        return {(au, ax, ay): 1}
    # x^ax y^ay = x^(ax-1) y^(ay-1) * (u - 1)
    rest = _reduce_one(au, ax - 1, ay - 1)
    out = {}
    for (au2, ax2, ay2), c in rest.items():
        # term with +u
        k1 = (au2 + 1, ax2, ay2)
        out[k1] = out.get(k1, 0) + c
        # term with -1
        k2 = (au2, ax2, ay2)
        out[k2] = out.get(k2, 0) - c
    # merge and drop zeros
    return {k: v for k, v in out.items() if v != 0}


def _add(A, B):
    out = dict(A)
    for k, v in B.items():
        out[k] = out.get(k, 0) + v
    return {k: v for k, v in out.items() if v != 0}


def _mul(A, B):
    out = {}
    for (au1, ax1, ay1), c1 in A.items():
        for (au2, ax2, ay2), c2 in B.items():
            prod = _reduce_one(au1 + au2, ax1 + ax2, ay1 + ay2)
            for k, v in prod.items():
                out[k] = out.get(k, 0) + c1 * c2 * v
    return {k: v for k, v in out.items() if v != 0}


def _poly_to_reduced(coeffs):
    """coeffs: {(ax,ay): c} -> reduced monomials {(au,ax,ay): c} (ax==0 or ay==0)."""
    out = {}
    for (ax, ay), c in coeffs.items():
        red = _reduce_one(0, ax, ay)
        for k, v in red.items():
            out[k] = out.get(k, 0) + c * v
    return {k: v for k, v in out.items() if v != 0}


def modular_bivariate(f_coeffs, e, m, t, X, Y):
    """
    Herrmann-May modular bivariate small-root finder (jvdsn port).
    f(x,y) as {(ax,ay): coeff}. Returns list of (x0, y0) roots.
    """
    U = X * Y

    def ev(au, ax, ay):
        return (U ** au) * (X ** ax) * (Y ** ay)

    f_red = _poly_to_reduced(f_coeffs)

    shifts = []
    for k in range(m + 1):
        for i in range(m - k + 1):
            # x^i * f^k * e^(m-k)
            acc = {(0, 0, 0): 1}
            for _ in range(k):
                acc = _mul(acc, f_red)
            # multiply by x^i
            acc2 = {}
            for (au, ax, ay), c in acc.items():
                r = _reduce_one(au, ax + i, ay)
                for kk, vv in r.items():
                    acc2[kk] = acc2.get(kk, 0) + c * vv
            scale = e ** (m - k)
            g = {(au, ax, ay): c * scale for (au, ax, ay), c in acc2.items()}
            shifts.append(g)
    for j in range(1, t + 1):
        for k in range((m // t) * j, m + 1):
            acc = {(0, 0, 0): 1}
            for _ in range(k):
                acc = _mul(acc, f_red)
            acc2 = {}
            for (au, ax, ay), c in acc.items():
                r = _reduce_one(au, ax, ay + j)
                for kk, vv in r.items():
                    acc2[kk] = acc2.get(kk, 0) + c * vv
            scale = e ** (m - k)
            g = {(au, ax, ay): c * scale for (au, ax, ay), c in acc2.items()}
            shifts.append(g)

    monomials = sorted({(au, ax, ay) for g in shifts for (au, ax, ay) in g})
    mono_idx = {mono: i for i, mono in enumerate(monomials)}
    nn = len(monomials)
    # square matrix: pad shifts with zero rows up to nn
    nr = len(shifts)
    L = IntegerMatrix(nn, nn)
    for row in range(nr):
        for (au, ax, ay), c in shifts[row].items():
            L[row, mono_idx[(au, ax, ay)]] = c * ev(au, ax, ay)

    L = LLL.reduction(L)

    # reconstruct polynomials h_i = sum L[row,col]*monomial / ev(bounds)
    polys = []
    for row in range(nn):
        poly = {}
        for col, (au, ax, ay) in enumerate(monomials):
            v = int(L[row, col])
            if v == 0:
                continue
            denom = ev(au, ax, ay)
            q = v // denom
            if q == 0:
                continue
            poly[(au, ax, ay)] = q
        if poly:
            polys.append(poly)

    # find roots: fast gcd method (jvdsn find_roots_gcd) — checks if
    # gcd(h1, h2) is a linear a*x + b*y  (no constant) -> roots (b,-a),(-b,a).
    x_s, y_s = sp.symbols('x y')

    def to_sympy(poly):
        expr = 0
        for (au, ax, ay), c in poly.items():
            for b in range(au + 1):
                cb = comb(au, b)
                expr += c * cb * (x_s ** (ax + b)) * (y_s ** (ay + au - b))
        return sp.Poly(expr, x_s, y_s, domain='ZZ')

    spolys = [to_sympy(p) for p in polys if p]
    if len(spolys) < 2:
        return []
    roots_found = []
    # pairwise gcd (jvdsn find_roots_gcd): a linear a*x+b*y -> roots (b,-a),(-b,a)
    found = False
    for i in range(len(spolys)):
        for j in range(i):
            g = sp.gcd(spolys[i], spolys[j])
            if g.total_degree() == 1 and len(g.gens) == 2:
                a = int(g.coeff_monomial(x_s))
                b = int(g.coeff_monomial(y_s))
                if (a, b) != (0, 0):
                    for (x0, y0) in [(b, -a), (-b, a)]:
                        if x0 != 0:
                            roots_found.append((x0, y0))
                    found = True
    if found:
        return roots_found
    # fallback: resultant (slower) of first two polys
    try:
        res = sp.Poly(sp.resultant(spolys[0], spolys[1], x_s), y_s, domain='ZZ')
        for yr in res.all_roots():
            if yr.is_integer:
                y0 = int(yr)
                up = sp.Poly(spolys[0].eval(y_s, y0), x_s, domain='ZZ')
                for xr in up.all_roots():
                    if xr.is_integer:
                        roots_found.append((int(xr), y0))
    except Exception:
        pass
    return roots_found


def boneh_durfee(e, n, delta=0.26, m=4):
    # jvdsn/crypto-attacks exact params (basic case, no partial p):
    # A = N + 1,  f = x*(A + y) + 1,  X = e^delta,  Y = 2^(n_bits/2 + 1)
    A = n + 1
    f_coeffs = {(1, 0): A, (1, 1): 1, (0, 0): 1}  # f = x*(A+y) + 1
    # jvdsn exact: X = e^delta, Y = 2^(n_bits/2 + 1)
    X = int(e ** delta)
    Y = int(2 ** (n.bit_length() // 2 + 1))
    t = int((1 - 2 * delta) * m)
    if t < 1:
        t = 1

    roots = modular_bivariate(f_coeffs, e, m, t, X, Y)
    for x0, y0 in roots:
        # recovery (jvdsn attack): require f(x0,y0) == 0 mod e
        z = x0 * (A + y0) + 1
        if z % e != 0:
            continue
        k = pow(x0, -1, e)
        s = (n + 1 + k) % e
        phi = n - s + 1
        ss = n - phi + 1
        disc = ss * ss - 4 * n
        if disc < 0:
            continue
        root = math.isqrt(disc)
        if root * root == disc:
            p = (ss + root) // 2
            q = (ss - root) // 2
            if p * q == n and isPrime(p) and isPrime(q):
                d = pow(e, -1, phi)
                return d
    return None


def main():
    pub = RSA.importKey(open(f"{HERE}/pub.key").read())
    e, n = pub.e, pub.n
    ct = int.from_bytes(open(f"{HERE}/flag.enc", "rb").read(), "big")
    print(f"n bits = {n.bit_length()}, e bits = {e.bit_length()}")

    d = boneh_durfee(e, n, delta=0.292, m=4)
    if d is None:
        print("Boneh-Durfee did not converge.")
        return None

    m = pow(ct, d, n)
    flag = long_to_bytes(m)
    if not flag.startswith(b"ASIS"):
        flag = b"\x00" + flag
    print("FLAG =", flag.decode(errors="replace"))
    return flag


if __name__ == "__main__":
    main()