# HITCON CTF 2022 - Easy NTRU
## prob.sage

from Crypto.Util.number import bytes_to_long as b2l
from Crypto.Util.number import long_to_bytes as l2b
import random

Zx.<x> = ZZ[]
def convolution(f,g):
  return (f * g) % (x^n-1)

def balancedmod(f,q):
  g = list(((f[i] + q//2) % q) - q//2 for i in range(n))
  return Zx(g)  % (x^n-1)

def randomdpoly(d1, d2):
  result = d1*[1]+d2*[-1]+(n-d1-d2)*[0]
  random.shuffle(result)
  return Zx(result)

def invertmodprime(f,p):
  T = Zx.change_ring(Integers(p)).quotient(x^n-1)
  return Zx(lift(1 / T(f)))

def invertmodpowerof2(f,q):
  assert q.is_power_of(2)
  g = invertmodprime(f,2)
  while True:
    r = balancedmod(convolution(g,f),q)
    if r == 1: return g
    g = balancedmod(convolution(g,2 - r),q)

def keypair():
  while True:
    try:
      f = randomdpoly(61, 60)
      f3 = invertmodprime(f,3)
      fq = invertmodpowerof2(f,q)
      break
    except Exception as e:
      pass
  g = randomdpoly(20, 20)
  publickey = balancedmod(3 * convolution(fq,g),q)
  secretkey = f
  return publickey, secretkey, g

def encode(val):
    poly = 0
    for i in range(n):
        poly += ((val%3)-1) * (x^i)
        val //= 3
    return poly

def encrypt(message, publickey):
  r = randomdpoly(18, 18)
  return balancedmod(convolution(publickey,r) + encode(message), q)


n, q = 263, 128
publickey, _, _ = keypair()

flag = open('flag', 'rb').read()
print(publickey)

for i in range(24):
    print(encrypt(b2l(flag), publickey))

## NTRU Cryptosystem

The question is implementing the NTRU Cryptosystem, and it encrypts the flag with the same public-key 24 times.

NTRUEncrypt procress is described here: https://en.wikipedia.org/wiki/NTRUEncrypt. And the encryptions are generated as such:

$$e_i = h*r_i + m \bmod q$$

This equation is under the Quotient Ring: $\mathbb {Z} _{128}[X]/(X^{263} - 1)$. $h$ is the public key, $m$ is the message (in this case the flag) and $r_i$ is a randomly generated polynomial with coeffcients in the set $\{-1, 0, 1\}$.

## Retransmission Attack

According this pdf: https://ntru.org/f/tr/tr006v1.pdf, if you encrypt the same message more than once with the same public key, then NTRU is vulnerable to the retransmission attack.

Using the differences in the encryptions, once can recover each individual term of the random polynomial $r_0$, which once fully recovered we can solve for $m$.

In other words if we take the difference between every encryption and the first one:

$$e_i - e_0 = h*(r_i - r_0) \bmod q\\
h^{-1}*(e_i - e_0) = r_i - r_0 \bmod q$$

As long as we can calculate $h^{-1}*(e_i - e_0)$, we will reveal the value of $r_i - r_0$, and with enough of these values we can recover the entirety of $r_0$.

## Calculating the public-key inverse

Unforunately the question isn't that easy, because $h^{-1}$ isn't defined in the Quotient Ring. The reason being because if you try to execute (from code given) `invertmodpowerof2(h)`, it will throw an exception saying that an inverse doesn't exist in $\mathbb {Z}_{2}[X]/(X^{263} - 1)$.

The way the code attempts to calculate an inverse in $\mathbb {Z} _{128}[X]/(X^{263} - 1)$ is by first calculating the inverse in $\mathbb {Z} _{2}[X]/(X^{263} - 1)$ and then using a reformation of Newton's formula for Hensel lifting to "lift" the inverse from $\mathbb {Z} _{2}$ to $\mathbb {Z} _{128}$. So the code fails to find an inverse in $\mathbb {Z}_{2}$, it obviously doesn't exist in $\mathbb {Z}_{128}$.

The reason the inverse doesn't exist is because $$X^{263} - 1 = (X - 1)*(X^{262} + X^{261} + X^{260} + \cdots + X^2 + X + 1)$$
In $Z_2$, $-1 = 1$, so the factorization is really $(X + 1)*(X^{262} + X^{261} + X^{260} + \cdots + X^2 + X + 1)$.

If you factor the public key under $\mathbb {Z}_{2}[X]$, you will see that it also contains a $(X + 1)$ factor, meaning that it shares a common factor with the modulus $X^{263} - 1$, and that's why an inverse doesn't exist.

However, that doesn't mean we can't solve for a "psuedo-inverse".

The modulus $X^{263} - 1$ isn't irreducible under $\mathbb {Z}_{2}[X]$, it decomposes into 3 factors. This means we decompose the Quotient Ring into the direct product of 3 Quotient Rings, where $(X - 1)*P_2*P_3 = X^{263} - 1$.

$$\mathbb{Z_2}[X]/(X^{263} - 1) \cong \mathbb{Z_2}[X]/(X - 1) \times \mathbb{Z_2}[X]/(P_2) \times \mathbb{Z_2}[X]/(P_3) $$

Now we can take the equation
$$e_i - e_0 = h*(r_i - r_0) \bmod q$$

And solve for $r_i - r_0$ over each of the individual Quotient Rings in $P_1, P_2$, and $P_3$, and then getting a solution in $\mathbb{Z_2}[X]/(X^{263} - 1)$ using the Chinese Remainder Theorem.

Then we can use a Hensel's lifting method to lift the solution to $\mathbb Z_{128}[X]/(X^{263} - 1)$.

We actually only need to lift twice to $\mathbb Z_{8}[X]/(X^{263} - 1)$, since we know the coefficients are in the set $\{-2, -1, 0, 1, 2\}$.

The good news is that $h^{-1}$ is well defined in $\mathbb{Z_2}[X]/(P_3)$ and $\mathbb{Z_2}[X]/(P_2)$. However it doesn't exist in $\mathbb{Z_2}[X]/(X - 1)$ (for the reasons described above).

We don't need to know the inverse to solve for what we need, though. The quotient ring $\mathbb{Z_2}[X]/(X - 1)$ is isomorphic to $Z_2$, because the only values possible are $0$ and $1$. Going back to our original equation, we basically have the following under that quotient ring:

$$e_i - e_0 = h*(r_i - r_0) \bmod 2\\
0 = 0*(r_i - r_0)$$

This means that $r_i - r_0$ is either $0$ or $1$ in $\mathbb{Z_2}[X]/(X - 1)$, so we can try both values when we calculate the CRT value, and get 2 choices for $r_i - r_0$ in $\mathbb{Z_2}[X]/(X^{263} - 1)$.

Onto Hensel lifting.

## Lifting from $\mathbb {Z_2}$ to $\mathbb Z_{128}$

Using a techinique similar to what was described here: https://math.stackexchange.com/questions/448946/lifting-linear-solutions-mod-p-to-mod-p2, we can do the following.

Using the above process, we can solve for $r_i - r_0$ in the equation below:
$$e_i - e_0 = h*(r_i - r_0) \bmod 2$$

To lift our solution for $r_i - r_0$ from $\bmod 2$ to $\bmod 4$, let $c_i$ be the solution to $r_i - r_0 \bmod 2$.
$$e_i - e_0 = h*c_i \bmod 2$$

Now we subsitute $r_i - r_0 = c_i + 2*d_i$:

$$e_i - e_0 = h*(c_i + 2*d_i) \bmod 4\\
e_i - e_0 - h*c_i = 2*h*d_i \bmod 4$$
We know that $e_i - e_0 - h*c_i = 0 \bmod 2$, so we can divide both sides of the equation by $2$.

$$\frac{e_i - e_0 - h*c_i}{2} = h*d_i \bmod 2\\
h^{-1}*\frac{e_i - e_0 - h*c_i}{2} = d_i \bmod 2$$

We can use the exact same process we did earlier to solve for $d_i \bmod 2$, as our definition $h^{-1}$ hasn't changed. We still solve for $d_i$ over the 3 Quotient Rings, and for the $\mathbb{Z_2}[X]/(X - 1)$ ring we just assume the solution is either $0$ or $1$.

After solving for $d_i$, we lift to $\bmod 8$. Let $r_i - r_0 = c_i + 2*d_i + 4*f_i$.
$$e_i - e_0 = h*(c_i + 2*d_i + 4*f_i) \bmod 8\\
\frac{e_i - e_0 - h*c_i - 2*h*d_i}{4} = h*f_i \bmod 8\\
h^{-1}*\frac{e_i - e_0 - h*c_i - 2*h*d_i}{4} = f_i \bmod 8$$

And now we have $r_i - r_0 = c_d + 2*d_i + 4*f_i \bmod 8$.

So in total we get 2 solutions in $\bmod 2$, 4 in $\bmod 4$, and 8 in $\bmod 8$. Only one of them is correct for the value of $r_i - r_0$, and it's very easy to check since the amount of terms will be very small compared to the rest.

Once we solve for $r_i - r_0$ we can proceed with the rest of the retransmission attack as described in the PDF.

## PoC

from Crypto.Util.number import bytes_to_long as b2l
from Crypto.Util.number import long_to_bytes as l2b
import random

Zx.<x> = ZZ[]
def convolution(f,g):
  return (f * g) % (x^n-1)

def balancedmod(f,q):
  g = list(((f[i] + q//2) % q) - q//2 for i in range(n))
  return Zx(g)  % (x^n-1)

def randomdpoly(d1, d2):
  result = d1*[1]+d2*[-1]+(n-d1-d2)*[0]
  random.shuffle(result)
  return Zx(result)

def invertmodprime(f,p):
  T = Zx.change_ring(Integers(p)).quotient(x^n-1)
  return Zx(lift(1 / T(f)))

def invertmodpowerof2(f,q):
  assert q.is_power_of(2)
  g = invertmodprime(f,2)
  while True:
    r = balancedmod(convolution(g,f),q)
    if r == 1: return g
    g = balancedmod(convolution(g,2 - r),q)

def keypair():
  while True:
    try:
      f = randomdpoly(61, 60)
      f3 = invertmodprime(f,3)
      fq = invertmodpowerof2(f,q)
      break
    except Exception as e:
      pass
  g = randomdpoly(20, 20)
  publickey = balancedmod(3 * convolution(fq,g),q)
  secretkey = f
  return publickey, secretkey, g

def encode(val):
    poly = 0
    for i in range(n):
        poly += ((val%3)-1) * (x^i)
        val //= 3
    return poly

def encrypt(message, publickey):
  r = randomdpoly(18, 18)
  return r, balancedmod(convolution(publickey,r) + encode(message), q)

n, q = 263, 128
publickey, _, _ = keypair()
flag = b'HITCON{test_flag_123141241}'

ris = []
eis = []

for _ in range(24):
    rr, ee = encrypt(b2l(flag), publickey)
    ris.append(rr)
    eis.append(ee)

h = publickey

def balancedmod_new(f,q,quo):
  g = list(((f[i] + q//2) % q) - q//2 for i in range(n))
  return Zx(g)  % quo

def convolution_new(f,g,quo):
  return (f * g) % quo

def invertmodprime_new(f,p,quo):
  T = Zx.change_ring(Integers(p)).quotient(quo)
  return Zx(lift(1 / T(f)))

def invertmodpowerof2_new(f,q,quo):
  #print(f, q, quo)
  assert q.is_power_of(2)
  g = invertmodprime_new(f,2,quo)
  #print("g:", g)
  while True:
    r = balancedmod_new(convolution_new(g,f,quo),q,quo)
    if r == 1: return g
    g = balancedmod_new(convolution_new(g,2 - r,quo),q,quo)

factors = [x^131 + x^129 + x^126 + x^123 + x^122 + x^118 + x^117 + x^116 + x^115 + x^114 + x^111 + x^109 + x^108 + x^107 + x^106 + x^104 + x^103 + x^102 + x^98 + x^97 + x^96 + x^95 + x^94 + x^93 + x^90 + x^89 + x^88 + x^86 + x^82 + x^81 + x^80 + x^78 + x^77 + x^76 + x^75 + x^73 + x^72 + x^71 + x^70 + x^67 + x^66 + x^63 + x^62 + x^61 + x^60 + x^59 + x^58 + x^57 + x^56 + x^53 + x^51 + x^48 + x^46 + x^39 + x^35 + x^32 + x^30 + x^29 + x^28 + x^25 + x^24 + x^22 + x^21 + x^19 + x^18 + x^17 + x^16 + x^15 + x^14 + x^13 + x^11 + x^10 + x^9 + x^8 + x^7 + x^6 + x^4 + x + 1, x^131 + x^130 + x^127 + x^125 + x^124 + x^123 + x^122 + x^121 + x^120 + x^118 + x^117 + x^116 + x^115 + x^114 + x^113 + x^112 + x^110 + x^109 + x^107 + x^106 + x^103 + x^102 + x^101 + x^99 + x^96 + x^92 + x^85 + x^83 + x^80 + x^78 + x^75 + x^74 + x^73 + x^72 + x^71 + x^70 + x^69 + x^68 + x^65 + x^64 + x^61 + x^60 + x^59 + x^58 + x^56 + x^55 + x^54 + x^53 + x^51 + x^50 + x^49 + x^45 + x^43 + x^42 + x^41 + x^38 + x^37 + x^36 + x^35 + x^34 + x^33 + x^29 + x^28 + x^27 + x^25 + x^24 + x^23 + x^22 + x^20 + x^17 + x^16 + x^15 + x^14 + x^13 + x^9 + x^8 + x^5 + x^2 + 1]
moduli = []
residues = []
for factor in factors:
    Ttmp = Zx.change_ring(Integers(2)).quotient(factor)
    h_part = Ttmp(h)
    moduli.append(Ttmp.modulus())
    residues.append(lift(Ttmp(invertmodprime_new(h, 2, factor))*Ttmp(eis[1] - eis[0])))

residues.append(0)
moduli.append(x + 1)
ci_cands = []
ci_cands.append(Zx(crt(residues, moduli))) # r_1 - r_0 in Z2/<x^n - 1>

# r_1 - r_0 can be either 0 or 1 mod (x + 1) we don't know
residues[-1] = 1
ci_cands.append(Zx(crt(residues, moduli)))

# now we lift from Z2 to Z4 then Z4 to Z8
#print("guess mod 2:", balancedmod(ci_cands[0], 2))
#print("guess mod 2:", balancedmod(ci_cands[1], 2))
#print("correct:", balancedmod(ris[1] - ris[0], 2))

cidi_cands = []
for ci_cand in ci_cands:
    xx3 = balancedmod(eis[1] - eis[0] - balancedmod(convolution(h, ci_cand), 4), 4)//2
    moduli = []
    residues = []
    for factor in factors:
        Ttmp = Zx.change_ring(Integers(2)).quotient(factor)
        h_part = Ttmp(h)
        moduli.append(Ttmp.modulus())
        residues.append(lift(Ttmp(invertmodprime_new(h, 2, factor))*Ttmp(xx3)))
    
    residues.append(0)
    moduli.append(x + 1)
    cidi_cands.append((ci_cand, Zx(crt(residues, moduli))))    
    
    residues[-1] = 1
    cidi_cands.append((ci_cand, Zx(crt(residues, moduli))))

#for ci_cand, di_cand in cidi_cands:
#    print("guess mod 4:", balancedmod(ci_cand + 2*di_cand, 4))
#print("correct:", balancedmod(ris[1] - ris[0], 4))

cidifi_cands = []
for ci_cand, di_cand in cidi_cands:
    xx5 = balancedmod(eis[1] - eis[0] - balancedmod(convolution(h, ci_cand), 8) - balancedmod(convolution(h, 2*di_cand), 8), 8) // 4
    moduli = []
    residues = []
    for factor in factors:
        Ttmp = Zx.change_ring(Integers(2)).quotient(factor)
        h_part = Ttmp(h)
        moduli.append(Ttmp.modulus())
        residues.append(lift(Ttmp(invertmodprime_new(h, 2, factor))*Ttmp(xx5)))
    
    residues.append(0)
    moduli.append(x + 1)
    cidifi_cands.append((ci_cand, di_cand, Zx(crt(residues, moduli))))    
    
    residues[-1] = 1
    cidifi_cands.append((ci_cand, di_cand, Zx(crt(residues, moduli))))

for ci_cand, di_cand, fi_cand in cidifi_cands:
    print("guess mod 8:", balancedmod(ci_cand + 2*di_cand + 4*fi_cand, 8))
print("correct:", balancedmod(ris[1] - ris[0], q))