Sharif CTF 8 Crypto 500+100pts: OSS hard Solver

solve.py

'''
Attacking to the OSS (Ong-Schnorr-Shamir) Signature Scheme: The perfect solution

References:
  [1] Jefferey Shallit. 1984. "An exposition of Pollard's algorithm for quadratic congruences"
  [2] Robin Chapman (https://math.stackexchange.com/users/226/robin-chapman), Efficiently finding two squares which sum to a prime, URL (version: 2010-10-06): https://math.stackexchange.com/q/5883
'''

from scryptos import *
from ecpy import modular_square_root
import gmpy
import math
import sys

debug = False

def absolutely_least_residue(a, m):
  '''
  Compute Absolutely Least Residue

  Args:
    a : An integer
    m : Modulus

  Returns:
    This function returns integer `x`, such that `x` \equiv `a` (mod `m`).
    Moreover, it satisfy the inequality "-`m`/2 <= `a` <= `m`/2".
  '''
  a %= m
  if 2*abs(a) > m:
    return a - m
  assert -m/2 <= a <= m/2
  return a

def find_uv(p, D):
  '''
  Find `u` and `v` which satisfy u^2 - Dv^2 \equiv 0\pmod p
    from the solution of x^2 - D \equiv 0\pmod p

  Args:
    p  : A prime
    D  : An integer

  Return:
    This function returns two-tuple: (u, v)
    There are the solution of this equation: u^2 - Dv^2 \equiv 0\pmod p
  '''
  C = gmpy.sqrt(abs(D))
  m0 = p
  x0 = min(*modular_square_root(D % p, p))
  xi = [x0]
  mi = [m0]

  i = 0
  while True:
    m_ = (xi[i]**2 - D) / mi[i]
    x_ = absolutely_least_residue(xi[i] % m_, m_)
    xi += [x_]
    mi += [m_]
    i += 1
    if mi[i] < C:
      break

  n = i
  Ai = {
    n-1: 1,
    n:0
  }
  for i in xrange(1, n):
    i = n-i-1
    Ai[i] = Ai[i+2] + ((xi[i+1] - xi[i]) / mi[i+1]) * Ai[i+1]

  u = abs(Ai[0] * xi[0]) - Ai[1] * mi[0]
  v = abs(Ai[0])
  assert (u**2 - D*v**2) % mi[0] == 0
  return (u, v)

def fermat_two_squares(p):
  '''
  Decompose the prime according to "Fermat's theorem on sums of two squares".

  Args:
    p : A prime which there exists `j` such that `p` = 4`j`+1.

  Returns:
    This function returns two integers-tuple `(x, y)`.
    These integer holds `x`^2 + `y`^2 = `p`.

  Note: 
    I quoted some functions (`powmods`, `quos`, `grem`, `ggcd`) and code from [2].
  '''
  def powmods(a, r, n):
    out = 1
    while r > 0:
      if (r % 2) == 1:
        r -= 1
        out = absolutely_least_residue(out * a, n)
      r /= 2
      a = absolutely_least_residue(a * a, n)
    return out

  def quos(a, n):
    if n <= 0:
      return "negative modulus"
    return (a - absolutely_least_residue(a, n))/n

  def grem(w, z):
    # remainder in Gaussian integers when dividing w by z
    (w0, w1) = w
    (z0, z1) = z
    n = z0 * z0 + z1 * z1
    assert n != 0, "division by zero"
    u0 = quos(w0 * z0 + w1 * z1, n)
    u1 = quos(w1 * z0 - w0 * z1, n)
    return(w0 - z0 * u0 + z1 * u1, w1 - z0 * u1 - z1 * u0)

  def ggcd(w, z):
    while z != (0,0):
      w, z = z, grem(w, z)
    return w

  k = gmpy.legendre(-1, p)
  if k == -1:
    print '-d is not quadratic residue'
    return None
  x0 = modular_square_root(-1 % p, p)[0]
  return ggcd((p, 0), (x0, 1))

def solve_two_squares_sum(d, k, n):
  '''
  Solve following equation:

  (1)    x^2 + d*y^2 \equiv k mod n

  Args:
    d : Plus/Minus flag (you can set only `+1` or `-1`).
    k : Target variable
    n : Modulus

  Returns:
    This function returns the solution of equation (1) as tuple of two integers.
  '''
  assert d in [1, -1], '`d` must be 1 or -1'
  if d == 1:
    p = k
    while True:
      if gmpy.is_prime(p) and (p - 1) % 4 == 0 and gmpy.legendre(-1, p) == 1:
        break
      p += n
    x, y = map(abs, fermat_two_squares(p))
    assert (x**2 + y**2) % n == k
    return x, y
  elif d == -1:
    if k % 2 == 1:
      r = k
    else:
      r = k + n
    x, y = (r+1)/2, (r-1)/2
    assert (x**2 - y**2) % n == k
    return x, y

def mult(x1, y1, x2, y2, D, n):
  '''
  Signature multiplication from `m1` and `m2`.

  Args:
    x1, y1 : Signature for `m1`
    x2, y2 : Signature for `m2`
    D, n   : Public parameter of that signature

  Returns:
    This function returns tuple of two-integers.
    It's signature corresponding to `m1`*`m2` mod `n`.
  '''
  return (
      (x1 * x2 + D * y1 * y2) % n,
      (y1 * x2 +     x1 * y2) % n
  )

def solve_weighted_two_squares_sum(n, D, k):
  '''
  Solve following equation:

  (1)    x^2 - Dy^2 \equiv k \pmod n

  Args:
    n : Modulus
    k : An integer - weight variable
    D : An integer - target variable

  Return:
    This function returns the solution of equation (1) as tuple of two integers.
  '''

  D = absolutely_least_residue(D, n)
  k = absolutely_least_residue(k, n)
  assert D != 0 and k != 0

  if gmpy.is_square(abs(D)) == 1:
    b = gmpy.sqrt(abs(D))
    if D >= 0:
      d = -1
    else:
      d = 1
    sol = solve_two_squares_sum(d, k, n)
    return (sol[0], sol[1] * modinv(b, n))

  if gcd(n, D) != 1 or gcd(n, k) != 1:
    print '[!] Factor found: %d' % max(gcd(n, D), gcd(n, k))
    0/0

  if abs(k) < abs(D):
    u, v = solve_weighted_two_squares_sum(n, k, D)
    return ((u * modinv(v, n)) % n, modinv(v, n))

  p = k % n
  while True:
    if gmpy.is_prime(p) and gmpy.legendre(D, p) == 1:
      break
    p += n
  u, v = find_uv(p, D)
  assert ((u**2 - D*v**2) % p) == 0
  _lambda = ((u**2 - D*v**2)/p) % n
  w, z = solve_weighted_two_squares_sum(n, D, _lambda)
  a = modinv(w**2 - D*z**2, n)
  s = (w*a) % n
  t = (z*a) % n
  x_sol, y_sol = mult(u, v, s, t, D, n)
  if debug:
    print '[+] D, k = %d, %d' % (D, k) 
    print '[+] Lambda = %d' % (_lambda)
    print '[+] Lambda*k = %d' % ((_lambda * p) % n)
    print '[+] u, v = %d, %d (u^2-Dv^2 = %d (mod n))' % (u, v, (u**2 - D*v**2) % n)
    print '[+] (u^2-Dv^2) * lambda^-1 = %d (mod n)' % ((((u**2 - D*v**2) % n) * modinv(_lambda, n)) % n)
    print '[+] (w^2-Dz^2)(s^2-Dt^2) = %d (mod n))' % ((((w**2 - D*z**2) % n) * (s**2 - D*t**2)) % n)
    print '[+] w, z = %d, %d (w^2-Dz^2 = %d (mod n))' % (w, z, (w**2 - D*z**2) % n)
    print '[+] x, y  = %d, %d  (x^2-Dy^2 = %d (mod n)' % (x_sol, y_sol, (x_sol**2 - D*y_sol**2) % n) 
    print '[+] Sanity Check: ', k, ((u**2 - D*v**2) % n * modinv(w**2 - D*z**2, n)) % n
    print '\n'
  assert (w**2 - D*z**2)%n == _lambda
  assert (x_sol**2 - D*y_sol**2) % n == k % n
  return x_sol, y_sol

def main():
  '''
  Sharif CTF 8 - OSS hard
  '''
  n = 108504503412454373447900779392896455789182908920252767679102572784833248190682207737154598741166656146295499891768592408962529857168166275863130804227243036676846199142906617137007877378696363607314850559328607693669984310189715960464177540321389050644453028598966277046571994650300388605135199566812211180551
  k = 96418402323876864512491639972930583881318685814783562179751981420110811774390896080021079024657494752001051891183358590923934531389848653055409236592228196204683162829359099715583891741617111941684481041179057606706801620321693724317614343290846702947941240579335556736392623484157714059707148044213280632330
  k = n-k
  public = (n, k)
  m = 9045418907335068770268779849124564950219260545189341826614770092040336019517687130751801431148236418688112027601673969058529643735762961447504773552645699712014099935774980274016297241177481291819364059206708658744657881196289540965565219550846730216720833999741328972855594202081834339846728109

  sys.setrecursionlimit(65535)
  x, y = solve_weighted_two_squares_sum(n, k, m)

  print x
  print y
  print (x**2 - k * y**2) % n == m

if __name__ == '__main__':
  main()