ASIS CTF 2018 Quals: Iran Solver (solved after contest finished)

solve.py

from scryptos import *
import itertools
import gmpy

pk_1 = '''-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQt5IN0XusElxNCZ3pj5YlTtqZiq
dORNqKaBilvYRmreqtHtX2MD5y3X3PShmF/eyC+Sz2EMHZVrWRNHW4mz7U1EcaDd
tVLQdnl1rB/wI+3dqr+3PTBAJ5uVQs/h7XxMjsxAFMxYFBmVze7gpyRRPXFPWP9K
4idj49HtS4rzKiRwXQIDAQAB
-----END PUBLIC KEY-----'''
pk_2 = '''-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAgPcib8EOESK34G9gAJx1
tapkjVSfe0E/7aiQtEIaamZv7JqBKucodEQP5w82sFECefpxj+aUfogURyyRAGxm
/t875h8RSHG97CUyJd1KYwR2NOzmfhAc/0jSNi61LZzaRuvpSNIvQowks2tLDEN8
Nh11Z0PNuz7xDygUHB9pG3Cc7ExAGv2hbr7gcUgAxxaiCnMIaVxF8sgRhkaV5nPB
yO2pQ2BjtTvHbsygC2JwaT/MdIwk04/yV4xJvRJOQOwR5M0mXEvepP5YzB9P1pZa
lkGGO9Y3RNZBjmmNzAWNQWxTuMtfHjnxM5LY295yXyVStuexi6tsiA+7oGys96pm
vQIDAQAB
-----END PUBLIC KEY-----'''
enc_0 = '''
BhCIhvtUQ63ItrnxPRFitkasBZMN1215gqvQb8QYA+bGyQiDrJnamT77HgsA3IkyEzow7E5YXQ+h
oZ9H7+YPeqlc03B0UUx2peJ5qXzAd9GTeG/5idghU8z46j0PfdZViD3d26PDlA4xrUXa3cY3IZpp
Iz9/Mk+G8CWchnjs2YPX'''.decode('base64')
enc_1 = '''
CO6jB7EkhdZ6tpIXTjj/kXXFPFA1fddECSV4rxrrsg0TLE/ptCn+e9lm9TBvodltY87C2uH3/mP1
evfIsaoxLzcX0tIPl16XCtuEYLkTQ+4S8vz7BQae5N7grivEUzh4T/EyF1Zf2s5zyDEZCorPXWj+
Xd0Qc2OwfoQ8+jpJ5znalsaTId0Z7ZI3z8BpWGF5bxVmvTOsNRf6G1PRYn9oI84WIaqbetW0OyoQ
BbGqzQApWk9fhbZVHxlaRk/+KouyQ1eLNZpvh1j7j2nJx3a7Lda4xP9/sdSrvKyVndYWZuKeZ1Bg
3uHuf/gQMa+o0CJrVUuTjMjMq1xFVTs1h1ooeA=='''.decode('base64')

pk_1 = RSA.import_pem(pk_1)
pk_2 = RSA.import_pem(pk_2)

print pk_1
print pk_2

# search rt, st by brute force t0, t1
for t0, t1 in itertools.product(xrange(1000), repeat=2):
  D = (4*t0 - t1)**2 + 16*pk_1.n
  d = is_perfect_square(D)
  if d != -1:
    y0 = (-(4*t0 + t1) + d)
    y1 = (-(4*t0 + t1) - d)
    if y0 % 8 == 0 or y1 % 8 == 0:
      ppq = max(y0 / 8, y1 / 8)
      if gmpy.is_prime(ppq + t0) == 0 or gmpy.is_prime(4*ppq + t1) == 0:
        continue
      rt, st = (ppq + t0), (4*ppq+t1)
      break
# rt = 22707497027074364781680987996518481005626032056396555613594154346975158102217036752539456409972400575111513400384893784030776676095511709063182297729868551
# st = 90829988108297459126723951986073924022504128225586222454376617387900632408868147010157825639889602300446053601539575136123106704382046836252729190919472251

assert rt * st == pk_1.n
print '[+] pk1.n = %d * %d' % (rt, st)

rps_cand = []

c = rt - 1
while True:
  if gmpy.next_prime(c) != rt:
    break
  if gmpy.next_prime(4*c) == st:
    rps_cand += [c]
  c -= 1

N = pk_2.n

for t in rps_cand:
  k = N - (t**2 + t + 1)
  d = is_perfect_square((t-1)**2 + 4*k)
  if d != -1:
    rps = t
    rs = (-(t - 1) + d) / 2
    d2 = is_perfect_square(rps**2 - 4*rs)
    assert d2 != -1
    r = (t + d2) / 2
    s = (t - d2) / 2
    break
print '[+] r = %d' % r
print '[+] s = %d' % s

p = gcd(r**3 - 1, N)
q = gcd(s**3 - 1, N)

print '[+] pk2.n = %d * %d' % (p, q)

assert r**2 + r + 1 == p
assert s**2 + s + 1 == q
assert p*q == N

res = long_to_bytes(RSA(pk_1.e, pk_1.n, p=rt, q=st).decrypt(bytes_to_long(enc_0)))
res += long_to_bytes(RSA(pk_2.e, pk_2.n, p=p, q=q).decrypt(bytes_to_long(enc_1)))
print [res]