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

# # SECCON 2019 CTF Quals - Crazy Repetition of Codes (crypto)
# 
# The challenge abbreviation is CRC, so we can guess what this challenge is about:
# 
# ```py
# import os
# from Crypto.Cipher import AES
# 
# def crc32(crc, data):
#   crc = 0xFFFFFFFF ^ crc
#   for c in data:
#     crc = crc ^ ord(c)
#     for i in range(8):
#       crc = (crc >> 1) ^ (0xEDB88320 * (crc & 1))
#   return 0xFFFFFFFF ^ crc
# 
# key = b""
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "TSG")
# assert(crc == 0xb09bc54f)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "is")
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "here")
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "at")
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "SECCON")
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000)):
#   crc = crc32(crc, "CTF!")
# key += crc.to_bytes(4, byteorder='big')
# 
# flag = os.environ['FLAG']
# assert(len(flag) == 32)
# 
# aes = AES.new(key, AES.MODE_ECB)
# encoded = aes.encrypt(flag)
# assert(encoded.hex() == '79833173d435b6c5d8aa08f790d6b0dc8c4ef525823d4ebdb0b4a8f2090ac81e')
# ```

# In order to decrypt the flag, we need to compute the standard $CRC32$ function on strings build by repeating a short chunk $11\ldots11$ times, where the number has 10000 **digits**! There is no hope to perform this number of operations (*at least* during the CTF). However, if we recall that $CRC32$ is an *affine* function, we can speed this up by using fast matrix exponentiation.
# 
# *In fact, the simplest solution is to recall that the order of $CRC32$ is $2^{32}-1$, and for any fixed chunk the order will divide this number. Therefore, we can simply reduce `int("1" x 10000)` modulo $2^{32}-1$, and run the code directly. After replacing the `crc32` implementation with `zlib.crc32`, it takes just a couple of minutes to compute the answer!*
# 
# <details>
# 
# ```py
# import os
# from Crypto.Cipher import AES
# from zlib import crc32
# 
# key = b""
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"TSG", crc)
# assert(crc == 0xb09bc54f)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"is", crc)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"here", crc)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"at", crc)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"SECCON", crc)
# key += crc.to_bytes(4, byteorder='big')
# 
# crc = 0
# for i in range(int("1" * 10000) % (2**32-1)):
#   crc = crc32(b"CTF!", crc)
# key += crc.to_bytes(4, byteorder='big')
# 
# flag = bytes.fromhex("79833173d435b6c5d8aa08f790d6b0dc8c4ef525823d4ebdb0b4a8f2090ac81e")
# aes = AES.new(key, AES.MODE_ECB)
# encoded = aes.decrypt(flag)
# print(repr(encoded))
# ```
# 
# </details>
# 
# <br/>
# 
# The idea is to construct the $32\times 32$ matrix $M$ and the 32-bit vector $c$ such that $CRC32_s(x, s) = M\times x \oplus c$ (where $s$ is the short string chunk used to build the string). This can be done in 33 black-box calls to $CRC32$: evaluating it at the zero vector and at all unit vectors. 
# 
# Finally, observe that the $e$-time repetition of $CRC32$ is expressed as
# $$
# CRC32^e(x, s) = M\times (\ldots M\times(M\times x \oplus c) \oplus c \ldots )\oplus c = M^e\times x \oplus (M^{e-1} \oplus M^{e-2} \oplus \ldots M \oplus I) \times c,
# $$
# where $I$ is the $32\times32$ identity matrix.
# 
# In our case $x=0$ so we only need to evaluate the sum $M^{e-1} \oplus M^{e-2} \oplus \ldots M \oplus I$. This can be done using standard geometric series formula (luckily, $M\oplus I$ is invertible):
# 
# $$
# (M^{e-1} \oplus M^{e-2} \oplus \ldots M \oplus I) = \frac{M^e\oplus  I}{M\oplus I}.
# $$
# 
# 

# In[1]:


from sage.all import *
from Crypto.Cipher import AES

def crc32(crc, data):
    crc = 0xFFFFFFFF ^^ crc
    for c in data:
        crc = crc ^^ ord(c)
        for i in range(8):
            crc = (crc >> 1) ^^ (0xEDB88320 * (crc & 1))
    return 0xFFFFFFFF ^^ crc

def frombin(v):
    return ZZ(tuple(map(int, v))[::-1], base=2)

def tobin(v, n):
    return tuple(ZZ(v).digits(2, padto=n))[::-1]

E = int("1" * 10000) % (2**32-1)
ID = identity_matrix(GF(2), 32, 32)

def rep(s):
    print("Computing for", repr(s))
    # build the vector
    c = crc32(0, s)

    # build the matrix
    m = matrix(GF(2), 32, 32)
    for i in range(32):
        v = [0] * 32
        v[i] = 1
        v = frombin(tuple(v))
        v = crc32(v, s) ^^ c
        v = tobin(v, 32)
        m.set_column(i, v)
    c = vector(GF(2), tobin(c, 32))

    matsum = (m**E + ID) * ~(m + ID)
    res = matsum * c
    return int(frombin(res)).to_bytes(4, byteorder="big")

key = b""
key += rep("TSG")
key += rep("is")
key += rep("here")
key += rep("at")
key += rep("SECCON")
key += rep("CTF!")

flag = bytes.fromhex("79833173d435b6c5d8aa08f790d6b0dc8c4ef525823d4ebdb0b4a8f2090ac81e")
aes = AES.new(key, AES.MODE_ECB)
encoded = aes.decrypt(flag)
print(encoded)