Google CTF 2021 Pythia

#!/usr/bin/python -u
import random
import string
import time

from base64 import b64encode, b64decode
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

max_queries = 150
query_delay = 10

passwords = [bytes(''.join(random.choice(string.ascii_lowercase) for _ in range(3)), 'UTF-8') for _ in range(3)]
flag = open("flag.txt", "rb").read()

def menu():
    print("What you wanna do?")
    print("1- Set key")
    print("2- Read flag")
    print("3- Decrypt text")
    print("4- Exit")
    try:
        return int(input(">>> "))
    except:
        return -1

print("Welcome!\n")

key_used = 0

for query in range(max_queries):
    option = menu()

    if option == 1:
        print("Which key you want to use [0-2]?")
        try:
            i = int(input(">>> "))
        except:
            i = -1
        if i >= 0 and i <= 2:
          key_used = i
        else:
          print("Please select a valid key.")
    elif option == 2:
        print("Password?")
        passwd = bytes(input(">>> "), 'UTF-8')

        print("Checking...")
        # Prevent bruteforce attacks...
        time.sleep(query_delay)
        if passwd == (passwords[0] + passwords[1] + passwords[2]):
            print("ACCESS GRANTED: " + flag.decode('UTF-8'))
        else:
            print("ACCESS DENIED!")
    elif option == 3:
        print("Send your ciphertext ")

        ct = input(">>> ")
        print("Decrypting...")
        # Prevent bruteforce attacks...
        time.sleep(query_delay)
        try:
            nonce, ciphertext = ct.split(",")
            nonce = b64decode(nonce)
            ciphertext = b64decode(ciphertext)
        except:
            print("ERROR: Ciphertext has invalid format. Must be of the form \"nonce,ciphertext\", where nonce and ciphertext are base64 strings.")
            continue

        kdf = Scrypt(salt=b'', length=16, n=2**4, r=8, p=1, backend=default_backend())
        key = kdf.derive(passwords[key_used])
        try:
            cipher = AESGCM(key)
            plaintext = cipher.decrypt(nonce, ciphertext, associated_data=None)
        except:
            print("ERROR: Decryption failed. Key was not correct.")
            continue

        print("Decryption successful")
    elif option == 4:
        print("Bye!")
        break
    else:
        print("Invalid option!")
    print("You have " + str(max_queries - query) + " trials left...\n")

from sage.all import *
from Crypto.Util.number import long_to_bytes
from Crypto.Cipher import AES
from bitstring import BitArray
import functools

P = PolynomialRing(GF(2), "x")
x = P.gen()
p = x ** 128 + x ** 7 + x ** 2 + x + 1
GH = GF(2 ** 128, "a", modulus=p)

def bytes_to_GH(data):
    """Simply convert bytes to field elements"""
    return GH([int(v) for v in BitArray(data).bin])

def GH_to_bytes(element):
    """Simply convert field elements to bytes"""
    return BitArray(element.polynomial().list()).tobytes().ljust(16, b'\x00')

def multi_collide_gcm(keyset, nonce, tag):
    R = PolynomialRing(GH, "r")
    L = bytes_to_GH(long_to_bytes(128 * len(keyset), 16))
    N = nonce + b'\x00\x00\x00\x01'
    T = bytes_to_GH(tag)
    interpolation_pts = []
    for key in keyset:
        H = bytes_to_GH(AES.new(key, AES.MODE_ECB).encrypt(b'\x00' * 16))
        B = ((L * H) + bytes_to_GH(AES.new(key, AES.MODE_ECB).encrypt(N)) + T) * H**-2
        interpolation_pts.append((H, B))
    sol = R.lagrange_polynomial(interpolation_pts)
    C_blocks = [GH_to_bytes(c) for c in sol.list()[::-1]]
    return b''.join(C_blocks) + tag

# cache the results for speedup, could have precomputed them but it's not that slow
@functools.lru_cache(maxsize=None)
def forge(start, end):
    keyset = list(keys.keys())[start:end]
    r = multi_collide_gcm(keyset, b'\x00'*12, b'\x01'*16)
    return r

import pickle
import itertools
import string
import base64
from sage.all import *
from Crypto.Util.number import long_to_bytes
from Crypto.Cipher import AES
from pwn import *
import functools
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.backends import default_backend
from bitstring import BitArray

def bytes_to_GH(data):
    """Simply convert bytes to field elements"""
    return GH([int(v) for v in BitArray(data).bin])

def GH_to_bytes(element):
    """Simply convert field elements to bytes"""
    return BitArray(element.polynomial().list()).tobytes().ljust(16, b'\x00')

def multi_collide_gcm(keyset, nonce, tag):
    R = PolynomialRing(GH, "r")
    L = bytes_to_GH(long_to_bytes(128 * len(keyset), 16))
    N = nonce + b'\x00\x00\x00\x01'
    T = bytes_to_GH(tag)
    interpolation_pts = []
    for key in keyset:
        H = bytes_to_GH(AES.new(key, AES.MODE_ECB).encrypt(b'\x00' * 16))
        B = ((L * H) + bytes_to_GH(AES.new(key, AES.MODE_ECB).encrypt(N)) + T) * H**-2
        interpolation_pts.append((H, B))
    sol = R.lagrange_polynomial(interpolation_pts)
    C_blocks = [GH_to_bytes(c) for c in sol.list()[::-1]]
    return b''.join(C_blocks) + tag

# cache the results for speedup, could have precomputed them but it's not that slow
@functools.lru_cache(maxsize=None)
def forge(start, end):
    keyset = list(keys.keys())[start:end]
    r = multi_collide_gcm(keyset, b'\x00'*12, b'\x01'*16)
    return r

def setKey(i):
    conn.sendline(b"1")
    conn.recvuntil(b">>> ")
    conn.sendline(f"{i}".encode())
    conn.recvuntil(b">>> ")

def decrypt(c):
    conn.sendline(b"3")
    conn.recvuntil(b">>> ")
    t = f"{nonce.decode()},{base64.b64encode(c).decode()}"
    conn.sendline(t.encode())
    conn.recvline()
    r = conn.recvline()
    conn.recvuntil(b">>> ")
    if b"Decryption failed." in r:
        return False
    return True

def bsearch(start, end):
    global tries
    mid = (end + start)//2
    if end - start == 1:
        return start
    tries -= 1
    print(f"tries left : {tries}")
    if decrypt(forge(start, mid)):
        return bsearch(start, mid)
    else:
        return bsearch(mid, end)

    
if __name__ == "__main__":
    # build all keys and pickle them for next time
    try:
        keys = pickle.load(open("keys.pickle", "rb"))
    except FileNotFoundError:
        keys = {}
        for t in itertools.product(string.ascii_lowercase, repeat=3):
            pwd = "".join(t).encode()
            print(pwd)
            kdf = Scrypt(salt=b'', length=16, n=2 ** 4, r=8, p=1, backend=default_backend())
            keys[kdf.derive(pwd)] = pwd
        pickle.dump(keys, open("keys.pickle", "wb"))

    # global variables
    P = PolynomialRing(GF(2), "x")
    x = P.gen()
    p = x ** 128 + x ** 7 + x ** 2 + x + 1
    GH = GF(2 ** 128, "a", modulus=p)

    tries = 150
    N = 26**3
    B = 500
    nonce = base64.b64encode(b'\x00' * 12)

    # recover the passwords and get the flag
    conn = remote("pythia.2021.ctfcompetition.com", 1337)
    # local testing
    # conn = process("./pythia.py")
    conn.recvuntil(b">>> ")

    password = b''
    # 3 passwords in total
    for j in range(3):
        # search in chunks
        for i in range(0, N, B):
            # if key is in this chunk
            if decrypt(forge(i, i + B)):
                print("Entering binary search...")
                index = bsearch(i, i+B)
                pwd = keys[list(keys.keys())[index]]
                password += pwd
                print(f"Found password : {pwd}")
                break
            tries -= 1
            print(f"tries left : {tries}")
        if j < 2:
            setKey(j+1)
            tries -= 1
            print(f"tries left : {tries}")
    print(f"full password = {password.decode()}")
    conn.sendline(b"2")
    conn.recvuntil(b">>> ")
    conn.sendline(password)
    conn.recvuntil(b"ACCESS GRANTED: ")
    print(f"Flag : {conn.recvline().decode()}")
    conn.close()