CODEGATE 2023 Finals - The Leakers (1 solve)

#!/usr/bin/sage
import random
import hashlib
import os
import signal 
 
signal.alarm(1800)
 
def PoW():
    prefix = os.urandom(8)
    print(prefix.hex())
    answer = bytes.fromhex(input().strip())
    assert len(answer) == 24
    result = hashlib.sha256(prefix + answer).digest()
    assert result[:3] == b"\x00\x00\x00"
 
P = PolynomialRing(ZZ, 'x')
x = P.gen()
 
def convolution(n, f, g):
    return (f * g) % (x ** n - 1)
 
def balance_mod(f, q):
    tt = f.coefficients(sparse = False)
    ret = 0
    for i in range(len(tt)):
        cc = int((tt[i] + q // 2) % q) - q // 2
        ret += cc * (x ** i)
    return ret
 
def random_poly(n, v1, v2):
    ret = v1 * [1] + v2 * [-1] + (n - v1 - v2) * [0]
    random.shuffle(ret)
    return P(ret)
 
def invert_prime(n, f, p):
    T = P.change_ring(GF(p)).quotient(x ** n - 1)
    ret = P(lift(1 / T(f)))
    return balance_mod(ret, 3)
 
def pad(n, arr):
    while len(arr) < n:
        arr.append(0)
    return arr
 
def encode(n, arr):
    res = 0
    for i in range(n):
        assert -1 <= arr[i] <= 1
        res += (arr[i] + 1) * (3 ** i)
    return res 
 
def task1(n, D):
    random.seed(int.from_bytes(os.urandom(32), "big"))
    f = random_poly(n, n // 3 + 1, n // 3)
    f3 = invert_prime(n, f, 3)
 
    random.seed(int.from_bytes(os.urandom(32), "big"))
    sel1 = random.sample(range(n), D)
    random.seed(int.from_bytes(os.urandom(32), "big"))
    sel2 = random.sample(range(n), D)
 
    coef_original = pad(n, f.coefficients(sparse = False))
    coef_inverse = pad(n, f3.coefficients(sparse = False))
 
    for i in range(D):
        coef_original[sel1[i]] = 0
        coef_inverse[sel2[i]] = 0
    
    print(sel1)
    print(sel2)
    print(encode(n, coef_original))
    print(encode(n, coef_inverse))
 
    assert int(input()) == encode(n, pad(n, f.coefficients(sparse = False)))
    assert int(input()) == encode(n, pad(n, f3.coefficients(sparse = False)))
 
def task2(n, D):
    random.seed(int.from_bytes(os.urandom(32), "big"))
    f = random_poly(n, n // 3 + 1, n // 3)
    f3 = invert_prime(n, f, 3)
    
    seed = int(input())
    random.seed(seed)
 
    sel1 = random.sample(range(n), D)
    sel2 = random.sample(range(n), D)
 
    coef_original = pad(n, f.coefficients(sparse = False))
    coef_inverse = pad(n, f3.coefficients(sparse = False))
 
    for i in range(D):
        coef_original[sel1[i]] = 0
        coef_inverse[sel2[i]] = 0
    
    print(sel1)
    print(sel2)
    print(encode(n, coef_original))
    print(encode(n, coef_inverse))
 
    assert int(input()) == encode(n, pad(n, f.coefficients(sparse = False)))
    assert int(input()) == encode(n, pad(n, f3.coefficients(sparse = False)))
 
PoW()
for _ in range(8):
    task1(2411, 83)
for _ in range(8):
    task2(8501, 2125)
 
flag = open("flag.txt", "r").read()
print(flag)

 

We are given two polynomials $f, f_v$ such that $f \cdot f_v \equiv 1 \pmod{x^n - 1}$, but some $D$ of the coefficients are erased. We have to recover $f, f_v$ completely, in a relatively fast and reliable fashion. The erasure positions are also given by the server.

 

For the first task, $(n, D) = (2411, 83)$ and the erasure positions are completely random. 

For the second task, $(n, D) = (8501, 2125)$ and the erasure positions can be controlled by a user provided seed. 

 

Task 1

By setting a variable for each erased coefficient, we will have a system of $n$ quadratic equations over $2D$ variables in $\mathbb{F}_3$. However, the interesting part is that some of the quadratic equations are actually just linear. For example, if we denote $S_1$ and $S_2$ as the set of erased coefficient's degree in $f$ and $f_v$ respectively, we can see that the equation arising from computing the coefficient of $x^k$ in $f \cdot f_v \pmod{x^n - 1}$ will be simply linear if there are no $u \in S_1$ and $v \in S_2$ such that $u + v \equiv k \pmod{n}$. 

 

By collecting these equations and solving the linear system, we will be closer to finding the solutions for the $2D$ variables.

However, after implementing this you can see that there will be a nontrivial kernel, of size around 40 to 50. 

 

This can be resolved in two ways. 

• The author's intended solution is to modify the given system as a system of $n$ quadratic equations over $K$ variables, where $K$ is the size of the kernel. This can be done simply by expressing the solution set of the $2D$ variables as a single solution added with a vector in the span of the computed kernel basis. As $K$ is much smaller than $2D$, we can actually solve this quadratic equation system by linearization. In other words, we can regard all quadratic terms as a separate linear variable, and solve the linear system over $\mathcal{O}(K^2)$ variables. This fails if $K$ is large, but such probability is small enough so that you can just try repeatedly.  
• soon-haari's solution works by selecting variables so that fixing it will add as many linear equations as possible, then brute forcing them. Apparently, brute forcing around 3 to 7 variables makes it sufficient to solve everything with a linear system. This was considered by the author as well, but was considered to be of similar difficulty. Congratulations to soon-haari for the solve!

 

Task 2

From solving task 1, it should be clear that the goal should be to create as many linear equations as possible, and the best way to do it is by making the erased coefficients consecutive in their positions. Note that $D = n/4$. Now how do we do that?

 

Looking at the sample implementation, we can see that the random sampling works by 

• selecting a random number below $n - i$
• taking the value at that index
• swapping with the value at position $n - i - 1$ so it's not selected again

if n <= setsize:
    # An n-length list is smaller than a k-length set.
    # Invariant:  non-selected at pool[0 : n-i]
    pool = list(population)
    for i in range(k):
        j = randbelow(n - i)
        result[i] = pool[j]
        pool[j] = pool[n - i - 1]  # move non-selected item into vacancy

 

The first idea is that our consecutive selections should be between $3n/4$ and $n$ - this is because if we try to pick everything from the front, the whole swapping process with the elements from the back makes everything very complicated. By picking everything at the back, the swapping process doesn't matter. Our goal is that for each $0 \le i < 2D$, the $i$th randbelow call should return a value $x$ such that $$n - D \le x < n - (i \pmod{D})$$ To do this efficiently, we need to minimize the number of bits we constrain from the randbelow results.

 

This can be done by finding $t, e$ such that $$n - D \le t \cdot 2^e < (t + 1) \cdot 2^e \le n - (i \pmod{D})$$ and maximizing $e$. Now, it suffices to constrain that the randbelow result is equal to $t$ after being shifted right $e$ bits. 

 

With this constraint in mind, finding the random seed is a relatively standard & well-known part. Check 

• https://gist.github.com/maple3142/1e3e81411a791f85073c3d902b0f14ef    
• https://rbtree.blog/posts/2021-05-18-breaking-python-random-module/    

 

from sage.all import * 
from pwn import * 
import os, time
from tqdm import tqdm
 
master_seed = 60024794038789808135493353697686616908493063225185396948889991711002019881012260998892909674537714109874172336760379200731611344495379266686257170624293028615555926994680615436128867836845735901814299032623034279895233846538952693992950207880205125875448087855836772541821353895968629329852908028951482292046429634872159757554409659745842900181726502284855318484066778535471595292549152348723264514710398644809539021921578114599851865111153331428784733149637717776626956886869295212600036738540930509437369722973201969837169465390106295804865432397396740692302797829400696350801197455057637221980617059842112641232969233787118965053967785222933031549323453023915303961246528225916509962869051725570255571786695958227355527242709785066976431806251198382777530714413831824130431722178073034419423922524991559359620024288713328086359403414933794337286671654435418265857597949112504278406015704720441696610732129686376522250690959070825600070450708353622473071494975779294350695405558303914530098365686912502766894493961489506186331053557290449360412540423892253756102487513806029093829893795865411528046974478581984437329446430455654183382959047805022370700975782191280373163986285947944114008525799059163331940037604447345980775790916325310887097685957858367094306509865450306900924968543630281114344776628409464077904562522114091218110199532185567329195460558546578726337659844292296095730138395687853467464686825377362140525018332064336498037333745669506410885086785171181512953442077146559888185321193265161347864107827758406549956890513622009827054983736589580372571827584944816635338237441722836404122046830295098128877031939911864423063671844869052268187006757023952046083313268542630040867977783706755657741531230633393603995284157154279750477826730638935203116032077204427304499208757288061746193707867350153416770764821598920213730711114965590743331727312353459097800387002408397658122900984840808246371821451013961389731787227009157398692371099660858182271036744448166340400895784346173287333244495255001791730081569336839520662381405144290409455974037142516983228317936699727842606081642240578844883094341186246580921199983915842081314638611291380152731368056730650469228559623662052531729887982086548223154333361363334651179104285034483322425275792671395416574805801066174116333550082050898468922185218775818163991247447213512072133729905831026979204387449740237107250035767499446745374366910625527910665062130250560706824973131871130316743119025554746764821434156674317934199946146598426673815806872877856328467376086085059620030855489151545303043307029412045887534894425887717583117437877318171620007894084301274965980741305242376340156346230267114570705189461331266648670135971622540649186222606301719344988665012528312715455324776650769063545393847491422947171743363094737848164203527815500665056795598783088145237029472339734427963849835022178538693065735805207641271044631695148529356782873975936571285135245429745146108036491430498162500763656245871098506267158087440237750159751664863211065697979948218429442546498813889729910223944717520161604727418685340138991981985619225849416335059973922841803883389393526433881970295102025473967271242147333800098274311358630381328801550779263853582690762954190377450436922016091335950930865947431088321016971794583463543506805777276891066135310863296397756985031896218023594355682462504506734940500365812292755553197773898889660699577561530145492927609494642174649101705826746672666183294279559291981779160964522297534960255063618332466976614501183228998682729197968575051178082566996628571473801354118883690941507931744401374987244372715844496021747653016503181871565513365490311937754665491041784007233799634349717657539855346128371782526748645254711640975576359645995419371902147356469195156372805778603944262661098290749863028288989422941985200606780481844930898792872531924601622394832468102957298372210031427732699319934781107988301247600593511823195668971916883951609735878765097692081037005879736118141874324737540398705286288988681436765122193427285686784758267811646782109984193537493517035025336405137273109568869777622232503118830613665850700496192371337485618910625319419657584661501620375974836832411786322507731977027575625783250266414378787148900895983949856500564564779297830591427575079297862649141244375485805801593463197364935921063290595220319427752540805812417257433276354627419044442018101969592668555978348415720913743417010710043519932066341337349002366227396881340998080393401376080508879416483875634105463190312058370838700500096751599290858489939320163312460446194287943741486626465044167916764309453265867227749001640757524516320238197501217997176854839242256247737322846205881431327708198117910578431481983383453719767475910167265170827225409678070198207927333097671611616080671517652067086557349582167623754660313732259984661354739151828000954072261669966765373367625577786880673420597700680012679943342785274712575950904993112141660047795983847702513444872812047642230003916360543076277750525218857399197096056760230787156929311820632995878912539801079019446915728253563406722059879474176881696819163446675139576322957663476859383520766444528469597928076769925748138660178853642739598139933326858259371267626938125840719239557084994647877357315029370116939357392564073911953692514212255988095253356674531690329478518006539330700869866952658773265849967085817409436103360317583749880904176992380298324586178020697936498010899061947345939204960967177164006648986295804937160615262813770879693445967104259390803516524422802908637470251797969593148039038477540080340954624195569529253847629154075274257199861147125055216614817003693519884916983205377401473762967508445819467333797312883675275497717519827470702088770584690428550650158813144401296875683289360897009594467643245700773535411913797121904382044029865707340324833895400600974009981197992647586507846704616021687940777933174900551897716451821657379278722485819954967117168842871510763631682184488735137792120697666469419973009451057094047507237995094164031163468702
 
P = PolynomialRing(ZZ, 'x')
x = P.gen()
 
# r = process(['sage', '../prob/for_organizer/task.sage'])
r = remote("52.79.59.27", 17776)
 
def solvePoW():
    PoW_start = time.time()
    prefix = bytes.fromhex(r.recvline().strip().decode())
    cnt = 0
    while True:
        cnt += 1
        if cnt % 1000000 == 0:
            print(cnt)
        ret = os.urandom(24)
        result = hashlib.sha256(prefix + ret).digest()
        if result[:3] == b"\x00\x00\x00":
            r.sendline(ret.hex())
            break
    PoW_end = time.time()
    print("PoW", PoW_end - PoW_start)
 
def balance_mod(f, q):
    tt = f.coefficients(sparse = False)
    ret = 0
    for i in range(len(tt)):
        cc = int((tt[i] + q // 2) % q) - q // 2
        ret += cc * (x ** i)
    return ret
 
def pad(n, arr):
    while len(arr) < n:
        arr.append(0)
    return arr
 
def encode(n, arr):
    res = 0
    for i in range(n):
        assert -1 <= arr[i] <= 1
        res += (arr[i] + 1) * (3 ** i)
    return res 
 
def decode(n, v):
    ret = [0] * n 
    for i in range(n):
        ret[i] = v % 3 - 1
        v = v // 3 
    return ret
 
def read_input(n, D):
    s1 = r.recvline()
    s2 = r.recvline()
    if b"Traceback" in s1:
        for _ in range(20):
            print(r.recvline())
    sel1 = eval(s1)
    sel2 = eval(s2)
    coef_original = decode(n, int(r.recvline().strip()))
    coef_inverse = decode(n, int(r.recvline().strip()))
    return sel1, sel2, coef_original, coef_inverse
 
def solve_init(n, D, sel1, sel2, coef_original, coef_inverse):
    mat, vec = [], []
    isok = [1 for _ in range(n)]
    for i in range(D):
        for j in range(D):
            isok[(sel1[i] + sel2[j]) % n] = 0
    
    idxs1 = [-1] * n 
    idxs2 = [-1] * n 
    for i in range(D):
        idxs1[sel1[i]] = i 
        idxs2[sel2[i]] = i
    
    for i in tqdm(range(n)):
        if isok[i] == 1:
            coefs = [0] * (2 * D)
            val = 0
            if i == 0:
                val = 1
            for j in range(n):
                if idxs1[j] != -1:
                    coefs[idxs1[j]] = coef_inverse[(i - j) % n]
                if idxs2[(i - j) % n] != -1:
                    coefs[idxs2[(i - j) % n] + D] = coef_original[j]
                if idxs1[j] == -1 and idxs2[(i - j) % n] == -1:
                    val -= coef_original[j] * coef_inverse[(i - j) % n]
            mat.append(coefs)
            vec.append(val % 3)
    
    M = Matrix(GF(3), mat)
    v = vector(GF(3), vec)
    sol = M.solve_right(v)
 
    kernel = M.right_kernel().basis()
    return sol, kernel 
 
def solve_task1(n, D):
    sel1, sel2, coef_original, coef_inverse = read_input(n, D)
 
    sol, kernel = solve_init(n, D, sel1, sel2, coef_original, coef_inverse)
 
    idxs1 = [-1] * n 
    idxs2 = [-1] * n 
    for i in range(D):
        idxs1[sel1[i]] = i 
        idxs2[sel2[i]] = i
 
    fvar = len(kernel)
    print(fvar)
    tot = fvar + (fvar * (fvar - 1)) // 2 + fvar
 
    idxs = [[0 for _ in range(fvar)] for _ in range(fvar)]
    for i in range(fvar):
        idxs[i][i] = i
    
    cur = fvar 
    for i in range(fvar):
        for j in range(i + 1, fvar):
            idxs[i][j] = cur
            cur += 1
    
    def single(x):
        return fvar + fvar * (fvar - 1) // 2 + x
 
    def wow(x1, x2):
        if x1 == x2:
            return x1
        else:
            if x1 > x2:
                x1, x2 = x2, x1
            return idxs[x1][x2]
 
    mat = []
    vec = []
 
    for i in tqdm(range(n)):
        coefs = [0] * tot
        val = 0
        if i == 0:
            val = 1
        for j in range(n):
            # [j] * [i - j]
            if idxs1[j] == -1 and idxs2[(i - j) % n] == -1:
                val -= coef_original[j] * coef_inverse[(i - j) % n]
            if idxs1[j] == -1 and idxs2[(i - j) % n] != -1:
                idx2 = idxs2[(i - j) % n]
                val -= coef_original[j] * sol[idx2 + D]
                for k in range(fvar):
                    coefs[single(k)] += coef_original[j] * kernel[k][idx2 + D]
            if idxs1[j] != -1 and idxs2[(i - j) % n] == -1:
                idx1 = idxs1[j]
                val -= coef_inverse[(i - j) % n] * sol[idx1]
                for k in range(fvar):
                    coefs[single(k)] += coef_inverse[(i - j) % n] * kernel[k][idx1]
            if idxs1[j] != -1 and idxs2[(i - j) % n] != -1:
                idx1 = idxs1[j]
                idx2 = idxs2[(i - j) % n]
                val -= sol[idx1] * sol[idx2 + D]
                for k in range(fvar):
                    coefs[single(k)] += sol[idx1] * kernel[k][idx2 + D]
                    coefs[single(k)] += sol[idx2 + D] * kernel[k][idx1]
                for k1 in range(fvar):
                    for k2 in range(fvar):
                        coefs[wow(k1, k2)] += kernel[k1][idx1] * kernel[k2][idx2 + D]
        mat.append(coefs)
        vec.append(val)
 
    M = Matrix(GF(3), mat)
    v = vector(GF(3), vec)
    final_sol = M.solve_right(v)
 
    fins = [0] * (2 * D)
 
    for i in range(2 * D):
        fins[i] += sol[i]
    
    for i in range(2 * D):
        for j in range(fvar):
            fins[i] += final_sol[single(j)] * kernel[j][i]    
 
    recover_f = 0
    recover_f3 = 0
    for i in range(n):
        if i in sel1:
            recover_f += fins[sel1.index(i)] * (x ** i)
        else:
            recover_f += coef_original[i] * (x ** i)
    
    for i in range(n):
        if i in sel2:
            recover_f3 += fins[sel2.index(i) + D] * (x ** i)
        else:
            recover_f3 += coef_inverse[i] * (x ** i)
    
    recover_f = balance_mod(recover_f, 3)
    recover_f3 = balance_mod(recover_f3, 3)
 
    r.sendline(str(encode(n, pad(n, recover_f.coefficients(sparse = False)))))
    r.sendline(str(encode(n, pad(n, recover_f3.coefficients(sparse = False)))))
 
def solve_task2(n, D):
    r.sendline(str(master_seed))
    sel1, sel2, coef_original, coef_inverse = read_input(n, D)
 
    for i in range(D):
        assert n - n // 4 <= sel1[i] < n 
        assert n - n // 4 <= sel2[i] < n
 
    sol, kernel = solve_init(n, D, sel1, sel2, coef_original, coef_inverse)
 
    print("task2 kernel", len(kernel))
 
    recover_f = 0
    recover_f3 = 0
    for i in range(n):
        if i in sel1:
            recover_f += sol[sel1.index(i)] * (x ** i)
        else:
            recover_f += coef_original[i] * (x ** i)
    
    for i in range(n):
        if i in sel2:
            recover_f3 += sol[sel2.index(i) + D] * (x ** i)
        else:
            recover_f3 += coef_inverse[i] * (x ** i)
 
    recover_f = balance_mod(recover_f, 3)
    recover_f3 = balance_mod(recover_f3, 3)
 
    r.sendline(str(encode(n, pad(n, recover_f.coefficients(sparse = False)))))
    r.sendline(str(encode(n, pad(n, recover_f3.coefficients(sparse = False)))))
 
st = time.time()
 
solvePoW()
 
for _ in tqdm(range(8)):
    solve_task1(2411, 83)
 
for _ in tqdm(range(8)):
    solve_task2(8501, 2125)
 
print(r.recvline())
 
en = time.time()
 
print(en - st)