WACON-Qualifier-2023/Crypto/PSS writeup

from Crypto.Util.number import *
import os
from hashlib import sha256
from tqdm import tqdm

def cascade_hash(msg, cnt, digest_len):
    assert digest_len <= 32
    msg = msg * 10
    for _ in range(cnt):
        msg = sha256(msg).digest()
    return msg[:digest_len]

def seed_to_permutation(seed):
    permutation = ''
    msg = seed + b"_shuffle"
    while len(permutation) < 16:
        msg = cascade_hash(msg, 777, 32)
        msg_hex = msg.hex()
        for c in msg_hex:
            if c not in permutation:
                permutation += c

    return permutation

def permutation_secret_sharing_gen(secret):
    seed_len = 5
    master_seed = os.urandom(seed_len)
    seed_tree = [None] * (2*N - 1)
    seed_tree[0] = master_seed
    for i in range(N-1):
        h = cascade_hash(seed_tree[i], 123, 2 * seed_len)
        seed_tree[2*i+1], seed_tree[2*i+2] = h[:seed_len], h[seed_len:]
    
    secret_list = list(secret.decode()) # ex) ['0','1','2','3',...]
    for i in range(N):
        # i-th party has a permutation derived from seed_tree[i+N-1]
        permutation = seed_to_permutation(seed_tree[i + N - 1])
        secret_list = [hex(permutation.find(x))[2:] for x in secret_list]

    permutated_secret = ''.join(secret_list)
    hidden_party = os.urandom(1)[0] & 7
    proof_idxs = merkle_proof_indexes[hidden_party]

    return seed_tree[proof_idxs[0]] + \
           seed_tree[proof_idxs[1]] + \
           seed_tree[proof_idxs[2]] + \
           bytes([hidden_party]) + \
           bytes.fromhex(permutated_secret)

merkle_proof_indexes = {
    0 : [2,4,8],
    1 : [2,4,7],
    2 : [2,3,10],
    3 : [2,3,9],
    4 : [1,6,12],
    5 : [1,6,11],
    6 : [1,5,14],
    7 : [1,5,13]
}


N = 8 # Number of parties

secret = b'---REDACTED---'
flag = b"WACON2023{" + secret + b'}'

assert len(secret) == 16 and set(secret) == set(b"0123456789abcdef")
# You can bruteforce the secret directly if you can overcome ^^^O(0xbeeeef * 16!)^^^!!!
assert cascade_hash(flag, 0xbeeeef, 32).hex() == 'f7a5108a576391671fe3231040777e9ac455d1bb8b84a16b09be1b2bac68345c'

fw = open("pss_data", "wb")
for _ in tqdm(range(2 ** 17)):
    fw.write(permutation_secret_sharing_gen(secret))

fw.close()

seed_len = 5
N = 8

merkle_proof_indexes = {
    0 : [2,4,8],
    1 : [2,4,7],
    2 : [2,3,10],
    3 : [2,3,9],
    4 : [1,6,12],
    5 : [1,6,11],
    6 : [1,5,14],
    7 : [1,5,13]
}
pss_data = open("pss_data", "rb").read()

def parse_single(base):
    hidden_party = pss_data[base + 3*seed_len]
    idx = merkle_proof_indexes[hidden_party]
    known_seed = [pss_data[base + i*seed_len: base + (i+1)*seed_len] for i in range(3)]
    permutated_secret = pss_data[base + 3*seed_len + 1: base + 3*seed_len + 1 + 8]
    return known_seed, idx, permutated_secret

def parse_data():
    known_seeds = {i:set() for i in range(3)}
    for i in range(2 ** 17):
        base = i * (3*seed_len + 1 + 8)
        known_seed, idx, _ = parse_single(base)
        known_seeds[idx[0]].add(known_seed[0])
    return known_seeds

known_seeds = parse_data()

import os
from hashlib import sha256

def cascade_hash(msg, cnt, digest_len):
    assert digest_len <= 32
    msg = msg * 10
    for _ in range(cnt):
        msg = sha256(msg).digest()
    return msg[:digest_len]

def seed_tree_gen(master_seed):
    seed_tree = [None] * (2*N - 1)
    seed_tree[0] = master_seed
    for i in range(N-1):
        h = cascade_hash(seed_tree[i], 123, 2 * seed_len)
        seed_tree[2*i+1], seed_tree[2*i+2] = h[:seed_len], h[seed_len:]
    return seed_tree

def check(master_seed, seed):
    base = pss_data.index(seed)
    known_seed, idx, _ = parse_single(base)
    seed_tree = seed_tree_gen(master_seed)
    return [seed_tree[idx[i]] == known_seed[i] for i in range(len(idx))] == [True, True, True]

def crack():
    master_seed = os.urandom(seed_len)
    h = cascade_hash(master_seed, 123, 2 * seed_len)
    h1, h2 = h[:seed_len], h[seed_len:]
    seed1 = None
    if h1 in known_seeds[1]:
        seed1 = h1
    if h2 in known_seeds[2]:
        seed1 = h2
    if seed1 is not None:        
        if check(master_seed, seed1):
            # print("master_seed:", master_seed, flush=True)
            return master_seed, seed1

from tqdm.auto import tqdm
import sys
import multiprocessing

def crack_(_):
    return crack()

ncpus = multiprocessing.cpu_count()
pool = multiprocessing.Pool(processes=ncpus-1)
iter_pool = pool.imap(crack_, range(2**23), chunksize=2*ncpus)

s = None
iter_tqdm = tqdm(iter_pool, total=2**23)
for r in iter_tqdm:
    if r is not None:
        iter_tqdm.close()
        print("Found, exiting brute.", flush=True)
        print("master_seed:", r[0], flush=True)
        pool.close()
        pool.terminate()
        s = r
        break

def seed2flag(master_seed, seed1):
    def seed_to_permutation(seed):
        permutation = ''
        msg = seed + b"_shuffle"
        while len(permutation) < 16:
            msg = cascade_hash(msg, 777, 32)
            msg_hex = msg.hex()
            for c in msg_hex:
                if c not in permutation:
                    permutation += c

        return permutation

    def recover_permutation(seed_tree, secret_list):
        for i in range(N-1, -1, -1):
            permutation = seed_to_permutation(seed_tree[i + N - 1])
            secret_list = [permutation[int(x, 16)] for x in secret_list]
        return secret_list

    seed_tree = seed_tree_gen(master_seed)
    base = pss_data.index(seed1)
    _, _, permutated_secret = parse_single(base)
    secret = recover_permutation(seed_tree, permutated_secret.hex())
    secret = ''.join(secret).encode()
    flag = b"WACON2023{" + secret + b'}'
    assert cascade_hash(flag, 0xbeeeef, 32).hex() == 'f7a5108a576391671fe3231040777e9ac455d1bb8b84a16b09be1b2bac68345c'
    print("flag:", flag)

if s is not None:
    seed2flag(*s)