import random
import sys
from hashlib import sha256
# Added bytes_to_long to the imports below
from Crypto.Util.number import long_to_bytes, bytes_to_long, inverse, GCD, isPrime

# --- Task Data ---
n = 104228256293611313959676852310116852553951496121352860038971098657350022997841589403091722735802150153734050783858816709247647536393314564077002364012463220999962114186339228164032217361145009468516448617173972835797623658266515762201804936729547278758839604969469770650218191574897316410254695420895895051693
phi = 104228256293611313959676852310116852553951496121352860038971098657350022997837434645707418205268240995284026522165519145773852565112344453740579163420312890001524537570675468046604347184376661743552799809753709321949095844960227307733389258381950812717245522599433727311919405966404418872873961877021696812800
q_dsa = 24513014442114004234202354110477737650785387286781126308169912007819
s1 = 764450933738974696530033347966845551587903750431946039815672438603
r1 = 8881880595434882344509893789458546908449907797285477983407324325035
s2 = 22099482232399385060035569388467035727015978742301259782677969649659

# r1 is used for both signatures (nonce reuse)
r2 = r1

def factor_using_phi(n, phi):
"""
Factors N given phi using a probabilistic method.
This works for any multi-prime N.
"""
print("[*] Starting factorization using phi...")

# We maintain a list of factors found so far (starting with N)
factors = [n]

# Calculate d and t such that phi = d * 2^t
t = 0
d = phi
while d % 2 == 0:
d //= 2
t += 1

final_primes = []

while len(factors) > 0:
curr = factors.pop()

# If the factor is prime, we are done with this branch
if isPrime(curr):
final_primes.append(curr)
continue

# Try to find a split for 'curr'
split_found = False
attempts = 0

while not split_found and attempts < 100:
attempts += 1
# Choose random base 'a'
a = random.randint(2, curr - 2)

# Compute x = a^d mod curr
x = pow(a, d, curr)

if x == 1 or x == curr - 1:
continue

# Repeatedly square x
for i in range(t - 1):
y = pow(x, 2, curr)
if y == curr - 1:
# Trivial square root path
break
if y == 1:
# Non-trivial square root of unity found!
# gcd(x - 1, curr) is a factor
factor = GCD(x - 1, curr)
factors.append(factor)
factors.append(curr // factor)
split_found = True
print(f" [+] Split found: {factor} ...")
break
x = y

return sorted(final_primes)

# 1. Factor N
primes = factor_using_phi(n, phi)
if len(primes) != 4:
print(f"[-] Error: Expected 4 primes, found {len(primes)}")
print(primes)
sys.exit(0)

print(f"[+] Recovered 4 primes: {primes}")

# 2. Reconstruct messages
# From task: n_factors = sorted(n_factors)
# m1 = n_factors[0] + n_factors[3]
# m2 = n_factors[1] + n_factors[2]
m1_int = primes[0] + primes[3]
m2_int = primes[1] + primes[2]

m1_bytes = long_to_bytes(m1_int)
m2_bytes = long_to_bytes(m2_int)

h1 = bytes_to_long(sha256(m1_bytes).digest())
h2 = bytes_to_long(sha256(m2_bytes).digest())

print(f"[+] Hash 1: {h1}")
print(f"[+] Hash 2: {h2}")

# 3. Exploit DSA Nonce Reuse
# s1 = k^-1 * (h1 + x*r) mod q
# s2 = k^-1 * (h2 + x*r) mod q
# Subtracting: k * (s1 - s2) = h1 - h2 mod q
# k = (h1 - h2) * (s1 - s2)^-1 mod q

try:
delta_h = (h1 - h2) % q_dsa
delta_s = (s1 - s2) % q_dsa

k = (delta_h * inverse(delta_s, q_dsa)) % q_dsa
print(f"[+] Recovered k: {k}")

# Recover private key x
# x = (s1 * k - h1) * r^-1 mod q
inv_r = inverse(r1, q_dsa)
numerator = (s1 * k - h1) % q_dsa
x = (numerator * inv_r) % q_dsa

print(f"[+] Recovered x: {x}")

flag = long_to_bytes(x)
print(f"\nFLAG: {flag.decode(errors='ignore')}")

except Exception as e:
print(f"[-] Error during DSA attack: {e}")