#!/usr/bin/env python
# coding: utf-8

# # CONFidence 2020 CTF Finals - FibHash (Crypto 421)

# In this challenge, we are given four linear recurrences (modulo a large prime $p$):
# $$
# \begin{align}
# (a, b) &\mapsto (6 a - 9 b, a),\\
# (a, b) &\mapsto (8 a - 15 b, a),\\
# (a, b) &\mapsto (10 a - 21 b, a),\\
# (a, b) &\mapsto (12 a - 35 b, a).
# \end{align}
# $$
# Each recurrence is applied $n$ times to some known starting values, where $n$ is the flag's integer value. We are given the first element of each output. The goal is as usual to recover the flag, i.e. find $n$.

# In[ ]:


# import gensafeprime
from Crypto.Util.number import bytes_to_long

flag = b"flag{local test flag}"
n = ZZ(bytes_to_long(flag))

# p = gensafeprime.generate(n.nbits() + 1)
# use challenge prime 
p = 3825410404746255934165193857151892506649946388303502923742559459860756754338954959067709726395434244171362767397769981378523965557840655385143008725493705117779476208930333229713993176785249759
assert n < p
print('p =', p)

K = Zmod(p)

def hash_with_params(g, h, a1, a0):
    def step(prev1, prev2):
        return vector([g * prev1 - h * prev2, prev1])
    def steps(n):
        Kx.<prev1, prev2> = K[]
        if n == 0: return vector([prev1, prev2])
        half = steps(n // 2)
        full = half(*half)
        if n % 2 == 0: return full
        else: return step(*full)
    return steps(n)[0](a1, a0)

starts = [
    (3141592653589793, 2384626433832795),
    (288419716939937, 5105820974944592),
    (3078164062862089, 9862803482534211),
    (7067982148086513, 2823066470938446),
]

# local test with sample flag
debug = 1
ys = [
    hash_with_params(6, 9,   3141592653589793, 2384626433832795),
    hash_with_params(8, 15,   288419716939937, 5105820974944592),
    hash_with_params(10, 21, 3078164062862089, 9862803482534211),
    hash_with_params(12, 35, 7067982148086513, 2823066470938446),
]


# Let's study the matrices $M_i$ of the steps of each linear recurrence. Jordan normal form is particularly useful in this case, as it will reveal a diagonalization if it exists. The hope is to find a basis change matrix $A_i$ such that $J_i = A_i^{-1} \times M_i \times A_i$ is diagonal matrix. Diagonal matrices are interesting because their $n$-th power is the diagonal matrix with each diagonal element raise to $n$'th power.

# In[2]:


Ms = []
Js = []
As = []
for g, h in [(6, 9), (8, 15), (10, 21), (12, 35)]:
    # decomposes over Z
    # otherwise, have to look in GF(p)
    M = matrix(ZZ, 2, 2, [
        g, -h,
        1, 0,
    ])
    J, A = M.jordan_form(subdivide=False, transformation=True)
    assert (~A) * M * A == J
    print(J, "\n")
    Ms.append(M.change_ring(K))
    Js.append(J.change_ring(K))
    As.append(A.change_ring(K))


# Interesting! We can see that the 2nd, 3rd and 4th matrices compute powers of 3, 5 and 7. For example:

# In[3]:


ys[1] == (As[1] * Js[1]**n * (~As[1]) * vector(starts[1]))[0] \
      == (As[1] * diagonal_matrix([K(5)**n, K(3)**n]) * (~As[1]) * vector(starts[1]))[0]


# We can also immediately notice that one side of the basis change simply modifies the starting points. However, we can not "remove" the other one, because we are given only one output of a pair. As a result, we obtain linear combinations of the diagonal powers:

# In[4]:


a0p, a1p = (~As[1]) * vector(starts[1])
a, b = As[1][0]
ys[1] == a * a0p * pow(5, n, p) + b * a1p * pow(3, n, p)


# It is likely that they are linearly independent and we can obtain values of $3^n,5^n,7^n$ separately. Let's try:

# In[5]:


coefs = []
for i in range(1, 4):
    a0p, a1p = (~As[i]) * vector(starts[i])
    a, b = As[i][0]
    coefs.append((a * a0p, b * a1p))
    
a, b = coefs[0]
c, d = coefs[1]
e, f = coefs[2]

# matrix mapping
# 3^n, 5^n, 7^n
# to outputs y1, y2, y3
m = matrix(GF(p), 3, 3, [
    b, a, 0,
    d, 0, c,
    0, f, e
])
m.rank()


# Great! We can obtain the powers now:

# In[6]:


assert (~m)[0] * vector(ys[1:4]) == pow(3, n, p)
assert (~m)[1] * vector(ys[1:4]) == pow(5, n, p)
assert (~m)[2] * vector(ys[1:4]) == pow(7, n, p)
select3 = (~m)[0]


# But what now? We have reduced the problem to standard discrete logarithm modulo $p$, which has 640 bits. This is rather infeasible for a CTF-level challenge.
# 
# But wait, we haven't used one of the outputs! Recall its Jordan form:

# In[7]:


Js[0]


# How do it's powers look like? Let's compute symbolically:

# In[8]:


_.<z> = QQ[]
print(matrix([[z, 1], [0, z]])**5)
print()
print(matrix([[z, 1], [0, z]])**10)


# Interesting! It seems to have terms $3^n$ and $n3^{n-1}$ only. Let's check for our large $n$:

# In[9]:


a0p, a1p = (~As[0]) * vector(starts[0])
a, b = As[0][0]
ys[0] == (a * a0p + b * a1p) * pow(3, n, p) + a * a1p * pow(3, n-1, p) * n 


# Cool! We also know $3^n$, so we can actually compute $n$ by simple arithmetics:

# In[10]:


pow3n = select3 * vector(ys[1:4])
coef0 = a * a0p + b * a1p
coef1 = a * a1p
nsol = (ys[0] - coef0 * pow3n) / (pow3n / 3) / coef1
nsol == n


# Now let's solve the challenge data. Note that all the coefficients are the same, we only change the $y$ values:

# In[11]:


ys = [
    1676692106337556305706016770958074706883705493218787835803500458160553522426125469104058304906236524961048084572483189178908787893713024829966510890627053066279418122862473691624486823188135802,
    2857964008499737244123466571337419134404523101896831014364752478262678705106912578400545946910842379193592617481759402862420715714781790030237246290673584587132334065412406222540095225886964680,
    377100011534086235658104379801079570118458049751793753941421646112566213186676075464534398000332714495426504495371864334278835331209624236632121660999291442470627222536838461859856916754666911,
    1452032184499983915936783392161922013770235576423037573850700970726680707755112154943113287273883628447114298114469426584835959464327248935513505491784021332818735403611670877092428106276909863,
]
b, c, d = -coef0 * select3
pow3n = select3 * vector(ys[1:4])
e, f, g = (K(coef1)/3) * select3
nsol = (ys[0] - coef0 * pow3n) / (pow3n / 3) / coef1
#nsol = (ys[0] + b * ys[1] + c * ys[2] + d * ys[3]) / (e * ys[1] + f * ys[2] + g * ys[3])
assert nsol * (e * ys[1] + f * ys[2] + g * ys[3]) ==  (ys[0] + b * ys[1] + c * ys[2] + d * ys[3])


# In[12]:


int(nsol).to_bytes(1000, "big").strip(b"\x00").decode()


# ## Bonus
# 
# Note that the solution is a degree-1 rational function of the outputs. Knowing (or guessing) this is sufficient for solving it in a black-box style! We can simply interpolate the polynomial $n\cdot g(y) = f(y)$, such that $n = f(y) / g(y)$:

# In[13]:


rows = []
for i in range(10):
    n = 100 + i
    ys = [
        hash_with_params(6, 9,   3141592653589793, 2384626433832795),
        hash_with_params(8, 15,   288419716939937, 5105820974944592),
        hash_with_params(10, 21, 3078164062862089, 9862803482534211),
        hash_with_params(12, 35, 7067982148086513, 2823066470938446),
    ]
    # for generality,
    # allow constants in numerator/denominator
    # but we don't need them in this problem
    row = []
    row += [1] + ys  # [1] for constants in numerator
    row += [n] + [n*y for y in ys] # [n] for constants in denominator
    rows.append(row)
m = matrix(K, rows)
print(m.nrows(), m.ncols(), m.rank())


# In[14]:


comb = m.right_kernel().matrix()[0]
top = -comb[:5]
bottom = comb[5:]


# In[15]:


ys = vector(K, [1] + [
    1676692106337556305706016770958074706883705493218787835803500458160553522426125469104058304906236524961048084572483189178908787893713024829966510890627053066279418122862473691624486823188135802,
    2857964008499737244123466571337419134404523101896831014364752478262678705106912578400545946910842379193592617481759402862420715714781790030237246290673584587132334065412406222540095225886964680,
    377100011534086235658104379801079570118458049751793753941421646112566213186676075464534398000332714495426504495371864334278835331209624236632121660999291442470627222536838461859856916754666911,
    1452032184499983915936783392161922013770235576423037573850700970726680707755112154943113287273883628447114298114469426584835959464327248935513505491784021332818735403611670877092428106276909863,
])
nsol = (top * ys) / (bottom * ys)


# In[16]:


int(nsol).to_bytes(1000, "big").strip(b"\x00")#.decode()