HITCON 2020 | another secret note

https://gist.github.com/samueltangz/bfc540af95a10e21a29e0f672ca048b8

#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

import os
from pwn import *
import requests
import base64
from Crypto.PublicKey import RSA
from Crypto.Cipher import AES
import hashlib
from gmpy2 import powmod
import json

from ctfools import Challenge as BaseChallenge
from ctfools import work # work(alg, check[, prefix, suffix, length, charset, threads])

class Challenge(BaseChallenge):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.solve_pow()

    def solve_pow(self):
        self.r.recvuntil(b'SHA256(XXXX+')
        suffix = self.r.recvuntil(b')')[:-1]
        self.r.recvuntil(b' == ')
        target = bytes.fromhex(self.r.recvline().strip().decode())
        charset = (string.ascii_letters+string.digits).encode()
        
        if not self.local:
            solution = work(hashlib.sha256, lambda h: h == target, suffix=suffix, length=4, charset=charset)
        else:
            solution = b'AAAA'

        self.r.sendline(solution)

    def register(self, name):
        self.r.sendlineafter(b'cmd: ', b'register')
        self.r.sendlineafter(b'name: ', name)
        self.r.recvuntil(b'token:  ')
        j = json.loads(base64.b64decode(self.r.recvline()))
        # AES-CBC encrypted (fixed key, unknown IV)
        return bytes.fromhex(j.get('cipher'))
        '''
        Example:
        Encrypts: b'{"secret": "hitcon{this_is_a", "who": "user", "name": "1234567"}\x10\x10...\x10'
        '''
    
    def login(self, ciphertext, iv=None):
        self.r.sendlineafter(b'cmd: ', b'login')
        j = {'cipher': ciphertext.hex()}
        if iv is not None:
            j['iv'] = iv.hex()
        token = base64.b64encode(json.dumps(j).encode())
        self.r.sendlineafter(b'token: ', token)
        
        self.r.recvuntil(B'token:  ')
        j = json.loads(base64.b64decode(self.r.recvline()))
        # AES-CBC encrypted (fixed key, unknown IV)
        return bytes.fromhex(j.get('cipher'))
        '''
        Do extra things if the JSON object cntains the 'cmd' entry. In particular:
        
        cmd == 'get_secret' and who == 'admin'  and name == 'admin'
        -> returns admin_secret
        cmd == 'get_time'
        -> returns time
        cmd == 'read_note' and has the field "note_name"
        -> returns note[note_name]
        cmd == 'note' and has the field "note"
        -> returns note_name and sets note[note_name] := note
        '''

def correct(local, plaintext, ciphertext):
    if not local: return True

    cipher = AES.new(b'YELLOW SUBMARINE', AES.MODE_ECB)
    return cipher.encrypt(plaintext) == ciphertext

# AES-CBC encryption of on{?????????", "
#                       ****************
def step1():
    local = 'local' in os.environ
    log = 'log' in os.environ

    r = Challenge(
        conn='nc 54.178.3.192 9427',
        proc=['python3', 'challenge/prob.py'],
        local=local,
        log=log)

    c = r.register(b'1234567')
    assert correct(local,
        xor(c[-32:-16], b'\x10'*16),
        c[-16:]
    )
    assert correct(local,
        xor(c[-48:-32], b'ame": "1234567"}'),
        c[-32:-16]
    )
    assert correct(local,
        xor(c[-64:-48], b'who": "user", "n'),
        c[-48:-32]
    )

    print(c[:32])

def recover(target_c, local, log, r, pos, charset=list(range(128))):
    if len(charset) == 1: return r, charset[0]
    ok = [32, 33] + list(range(35, 91+1)) + list(range(93, 127+1))

    best_c = 0
    best_d = 128

    # pick the best candidate to separate the charset into two (evenly)
    for c in range(128):
        charset_l, charset_r = [], []

        for k in charset:
            if c^k in ok:
                charset_l.append(k)
            else:
                charset_r.append(k)
        
        d = abs(len(charset_l) - len(charset_r))
        if d < best_d:
            best_c, best_d = c, d
    
    c = best_c
    charset_l, charset_r = [], []
    for k in charset:
        if c^k in ok:
            charset_l.append(k)
        else:
            charset_r.append(k)
        
    # sends the query
    if r is None:
        r = Challenge(
            conn='nc 54.178.3.192 9427',
            proc=['python3', 'challenge/prob.py'],
            local=local,
            log=log)

    forged_iv = xor(
        target_c[:16],
        b'on{?????????", "',
        b'["{?????????"]\2\2',
        b'\0'*(3+pos) + bytes([c]) + b'\0'*(12-pos)
    )
    forged_c = target_c[16:]

    print('c =', c, charset)

    good = False
    try:
        r.login(forged_c, iv=forged_iv)
        good = True
    except:
        r.r.close()
        r = None

    if good:
        return recover(target_c, local, log, r, pos, charset_l)
    else:
        return recover(target_c, local, log, r, pos, charset_r)

# Recover the unknowns in on{?????????", " -- hence the user secret
#                         ****************
def step2():
    local = 'local' in os.environ
    log = 'log' in os.environ

    if local:
        target_c = b':?\xf8zxm\xd8WLv\xb1\x94\x830\xa9\x1f`Y\xb1[Ry\x1a|V\x94\x0el%\xc4Y\xa8'
    else:
        target_c = b"\xb6\xf6\xc3\xc0J\xae\x166\xfb'$0\xbf\x7f\xdd\r\x14\x9dW\xb0\xc2\x06\xa5\x1c\xe5'Y\x11E\x1e\x15\xc5"

    r = None

    user_secret = b'hitcon{'
    for i in range(9):
        r, fc = recover(target_c, local, log, r, i)
        user_secret += bytes([fc])

    if local:
        assert user_secret == b'hitcon{this_is_a'
    else:
        print('user_secret =', user_secret)
        # hitcon{JSON_is_5

# Constructing {"cmd":"get_secret", "who": "admin", "name": "admin"}
# Returns      {"cmd": "get_secret", "who": "admin", "name": "admin", "secret": "_test_flag!!!!!}"}
def step3():
    local = 'local' in os.environ
    log = 'log' in os.environ

    if local:
        target_c = b':?\xf8zxm\xd8WLv\xb1\x94\x830\xa9\x1f`Y\xb1[Ry\x1a|V\x94\x0el%\xc4Y\xa8'
        ref_m = b'on{this_is_a", "'
    else:
        target_c = b"\xb6\xf6\xc3\xc0J\xae\x166\xfb'$0\xbf\x7f\xdd\r\x14\x9dW\xb0\xc2\x06\xa5\x1c\xe5'Y\x11E\x1e\x15\xc5"
        ref_m = b'on{JSON_is_5", "'


    r = Challenge(
        conn='nc 54.178.3.192 9427',
        proc=['python3', 'challenge/prob.py'],
        local=local,
        log=log)


    forged_iv = xor(target_c[:16], ref_m, b'{"ame":"admin"}\1')
    forged_ct = target_c[16:]
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"ame": "admin"}',
        b'{"name":"admin"}'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"name": "admin"',
        b'{"xname":"admin"'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xname": "admin',
        b'{"xxname":"admin'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxname": "admi',
        b'{"xxxname":"admi'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxname": "adm',
        b'{"xxxxname":"adm'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxxname": "ad',
        b'{"xxxxxname":"ad'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxxxname": "a',
        b'{"xxxxxxname":"a'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxxxxname": "',
        b'{"xxxxxxxname":"'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxxxxxname": ',
        b'{"xxxxxxxxname":'
    )
    c = r.login(forged_ct, forged_iv)
    # Checkpoint: m = b'{"xxxxxxxxname": "admin"}'

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xxxxxxxxname":',
        b'{"a":"n","name":'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"a": "n", "name',
        b'{"aa":"in","name'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"aa": "in", "na',
        b'{"aa":"dmin","na'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"aa": "dmin", "',
        b'{"who":"admin","'
    )
    c = r.login(forged_ct, forged_iv)
    # Checkpoint: m = b'{"who": "admin", "name": "admin"}'

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"who": "admin",',
        b'{"x":1,"a":"in",'
    )
    c = r.login(forged_ct, forged_iv)
    # Checkpoint: m = b'{"x": 1, "a": "in", "name": "admin"}'

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": 1, "a": "i',
        b'{"x":1,"a":"admi'
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": 1, "a": "a',
        b'{"xx":1,"who":"a',
    )
    c = r.login(forged_ct, forged_iv)
    # Checkpoint: m = b'{"xx": 1, "who": "admin", "name": "admin"}'

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"xx": 1, "who":',
        b'{"x":"et","who":',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "et", "who',
        b'{"x":"cret","who',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "cret", "w',
        b'{"x":"secret","w',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "secret", ',
        b'{"x":"t_secret",',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "t_secret"',
        b'{"x":"et_secret"',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "et_secret',
        b'{"x":"get_secret',
    )
    c = r.login(forged_ct, forged_iv)
    # Checkpoint: {"x": "get_secret", "who": "admin", "name": "admin"}

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"x": "get_secre',
        b'{"md":"get_secre',
    )
    c = r.login(forged_ct, forged_iv)

    forged_ct = c
    forged_iv = xor(
        forged_iv,
        b'{"md": "get_secr',
        b'{"cmd":"get_secr',
    )
    c = r.login(forged_ct, forged_iv)

    print(c)
    # {"cmd": "get_secret", "who": "admin", "name": "admin", "secret": "_test_flag!!!!!}"}

def recover2(target_c, local, log, r, pos, partial_admin_secret, charset=list(range(128))):
    if len(charset) == 1: return r, charset[0]
    ok = [32, 33] + list(range(35, 91+1)) + list(range(93, 127+1))

    best_c = 0
    best_d = 128

    # pick the best candidate to separate the charset into two (evenly)
    for c in range(128):
        charset_l, charset_r = [], []

        for k in charset:
            if c^k in ok:
                charset_l.append(k)
            else:
                charset_r.append(k)
        
        d = abs(len(charset_l) - len(charset_r))
        if d < best_d:
            best_c, best_d = c, d
    
    c = best_c
    charset_l, charset_r = [], []
    for k in charset:
        if c^k in ok:
            charset_l.append(k)
        else:
            charset_r.append(k)
        
    # sends the query
    if r is None:
        r = Challenge(
            conn='nc 54.178.3.192 9427',
            proc=['python3', 'challenge/prob.py'],
            local=local,
            log=log)

    forged_iv = xor(
        target_c[48:64],
        b' "????????????' + partial_admin_secret,
        b' "????????????"\1',
        b'\0'*(2+pos) + bytes([c]) + b'\0'*(13-pos)
    )
    forged_c = target_c[64:80]

    print('c =', c, charset)

    good = False
    try:
        r.login(forged_c, iv=forged_iv)
        good = True
    except:
        r.r.close()
        r = None

    if good:
        return recover2(target_c, local, log, r, pos, partial_admin_secret, charset_l)
    else:
        return recover2(target_c, local, log, r, pos, partial_admin_secret, charset_r)

def recover3(target_c, local, log, r, pos, partial_admin_secret, charset=list(range(128))):
    if len(charset) == 1: return r, charset[0]
    ok = [32, 33] + list(range(35, 91+1)) + list(range(93, 127+1))

    best_c = 0
    best_d = 128

    # pick the best candidate to separate the charset into two (evenly)
    for c in range(128):
        charset_l, charset_r = [], []

        for k in charset:
            if c^k in ok:
                charset_l.append(k)
            else:
                charset_r.append(k)
        
        d = abs(len(charset_l) - len(charset_r))
        if d < best_d:
            best_c, best_d = c, d
    
    c = best_c
    charset_l, charset_r = [], []
    for k in charset:
        if c^k in ok:
            charset_l.append(k)
        else:
            charset_r.append(k)
        
    # sends the query
    if r is None:
        r = Challenge(
            conn='nc 54.178.3.192 9427',
            proc=['python3', 'challenge/prob.py'],
            local=local,
            log=log)

    forged_iv = xor(
        target_c[16:32],
        partial_admin_secret + b'??}"}' + b'\n'*10, # known
        b'"??}" ' + b'\n'*10, # target
        b'\0'*(1+pos) + bytes([c]) + b'\0'*(14-pos)
    )
    forged_c = target_c[32:]

    print('c =', c, charset)

    good = False
    try:
        r.login(forged_c, iv=forged_iv)
        good = True
    except:
        r.r.close()
        r = None

    if good:
        return recover3(target_c, local, log, r, pos, partial_admin_secret, charset_l)
    else:
        return recover3(target_c, local, log, r, pos, partial_admin_secret, charset_r)

def step4():
    local = 'local' in os.environ
    log = 'log' in os.environ

    if local:
        target_c = b'[\x0c@\nw\xb7\x10\x84q\xf0\xcc\xf1\xff\x08\xfa\xae\xe9g3\xa1\xbf\xb6\xd4\x10v#\xb8\x9c\xa4Gd\xce2d\x92M\x1b0\xe1o\xf6\xfc\x94\x9a\x88mm/J^z\xdaS\xe8\xe2\x05\x8cm\xd6o\xc3g\x81\xe0\x99\x9ah\x8d!1r?q\x1e\xbb\x18\xbd%\xef\xa3\xa6~\x96\xaa\x90\x1d\xcd\x04?K\xd7\xca\r\x8f5B'
    else:
        target_c = b'V\xf4\xd6\xb6\xa7B\xbb\x8f\xaa`,\x83\x19\xe2\xa9\xdfs\xa0\xa48}\xa9JLN\xf2\xcbk\x7f\xfa\x82\xb0f\xa1\xbeL\xcf\xd6&u\xcab\r\xd7\x03\xd4</\x8c\xa9\xb2\xa0\x90\x9f\x8fG\xb1$\xd8Y\xf9\x11&\x19h\xf1\xaf=\xaf\x8f\xb5p\xdc\xc3\xbcJ_*Se\x9cD\xf9\xe8]\x14\xfdQY\xdc]\xa6n\x98\xafY'

    admin_secret = [None for _ in range(16)]

    # use existing record
    admin_secret[15] = 125
    # if not local:
    #     for i in range(14):
    #         admin_secret[i] = [48, 95, 119, 111, 78, 100, 99, 114, 70, 117, 108, 33, 64, 35][i]

    r = Challenge(
        conn='nc 54.178.3.192 9427',
        proc=['python3', 'challenge/prob.py'],
        local=local,
        log=log)

    forged_c = target_c[48:]
    forged_iv = xor(
        target_c[32:48],
        b'dmin", "secret":',
        b'{"a":[0,0],"c": '
    )    
    semi_c = r.login(forged_c, iv=forged_iv)
    # {"a": [0, 0], "c": "_test_flag!!!!!}"}
    
    # Recover the starred part below (13th byte):
    # {"cmd": "get_secret", "who": "admin", "name": "admin", "secret": "???????????????}"}
    #                                                                               *
    if admin_secret[12] is None:
        for i in range(32, 127+1):
            if r is None:
                r = Challenge(
                    conn='nc 54.178.3.192 9427',
                    proc=['python3', 'challenge/prob.py'],
                    local=local,
                    log=log)
            
            forged_c = semi_c[32:]
            forged_iv = xor(
                semi_c[16:32],
                bytes([i]) + b'??}"}' + b'\n'*10, # guess
                b'"??}" ' + b'\n'*10 # target
            )
            try:
                r.login(forged_c, iv=forged_iv)
                admin_secret[12] = i
                break
            except:
                r.r.close()
                r = None

    # the 14&15th byte
    for i in range(2):
        if admin_secret[13+i] is not None: continue
        r, fc = recover3(semi_c, local, log, r, i, bytes([admin_secret[12]]))
        admin_secret[13+i] = fc
        print(admin_secret) # DEBUG

    # First 12 bytes
    partial_admin_secret = bytes(admin_secret[12:13+1])
    for i in range(12):
        if admin_secret[i] is not None: continue
        r, fc = recover2(target_c, local, log, r, i, partial_admin_secret)
        admin_secret[i] = fc
        print(admin_secret) # DEBUG

    print(admin_secret)
    admin_secret = bytes(admin_secret)

    if local:
        assert admin_secret == b'_test_flag!!!!!}'
    else:
        print('admin_secret =', admin_secret)

if __name__ == '__main__':
    # step1()
    # step2() # hitcon{JSON_is_5
    # step3()
    step4() # 0_woNdcrFul!@##}

'''
ok = []
for x in range(128):
    try:
        json.loads('"' + chr(x) + '"')
        ok.append(x)
    except:
        pass
'''
# hitcon{JSON_is_50_woNderFul!@##}