BsBsこうしょう

これは考えたことではなく思ったことです。

nanoDiamond[writeup][WMCTF2022]

本番中では解けなかったが、元ネタの論文を見て面白かったので復習。

問題

# from Crypto.Util.number import *
import string
import secrets
from hashlib import sha256
from random import randint, shuffle, choice

def proof_of_work():
    s = ''.join([secrets.choice(string.digits + string.ascii_letters)
                for _ in range(20)])
    print(f'sha256(XXXX+{s[4:]}) == {sha256(s.encode()).hexdigest()}')
    if input('Give me XXXX: ') != s[:4]:
        exit(1)

ROUND_NUM = 50
PREROUND_NUM = 14
CHEST_NUM = 6

flag = 'FLAG{testtesttest}'
# with open('flag', 'r') as f:
#     flag = f.read()

white_list = ['==','(',')','0','1','and','or','B0','B1','B2','B3','B4','B5']

def calc(ans, chests, expr):
    B0, B1, B2, B3, B4, B5 = chests
    return ans(eval(expr))

def round():
    chests = [choice((True, False)) for _ in range(CHEST_NUM)]
    print("Six chests lie here, with mimics or treasure hidden inside.")
    print("But don't worry. Skeleton Merchant knows what to do.")
    print("Be careful, Skeleton Merchant can lie twice!")

    truth = lambda r: not not r
    lie = lambda r: not r
    lie_num = randint(0, 2)
    lie_status = [truth] * (PREROUND_NUM - lie_num) + [lie] * lie_num
    shuffle(lie_status)

    for i in range(PREROUND_NUM):
        try:
            question = input('Question: ').strip()
            for word in question.split(' '):
                assert word in white_list, f"({word}) No treasure for dirty hacker!"
            result = calc(lie_status[i], chests, question)
            print(f'Answer: {result}!')
        except Exception as e:
            print("Skeleton Merchant fails to understand your words.")
            print(e)
    print('Now open the chests:')
    return chests == list(map(int, input().strip().split(' ')))


if __name__ == '__main__':

    # proof_of_work()

    print('Terraria is a land of adventure! A land of mystery!')
    print('Can you get all the treasure without losing your head?')

    for i in range(ROUND_NUM):
        if not round():
            print('A chest suddenly comes alive and BITE YOUR HEAD OFF.')
            exit(0)
        else:
            print('You take all the treasure safe and sound. Head to the next vault!')

    print(f"You've found all the treasure! {flag}")

問題の本質に関係ないPoWの部分はコメントアウトしている。これをまとめると、以下のような問題になる。

  • \( 0 \leq e < 2 ^ 6 \) の整数 $ e $ を yes/no で答えられるクエリだけを使って的中させる
  • クエリの回数は14回
  • ただし、クエリの返答は0~2回嘘をつくことがある (何回かはランダム)

元ネタ

これは "Ulam's Searching Game With two Lies" という論文が元ネタになっている。

en.wikipedia.org

pdf.sciencedirectassets.com

この論文のほとんどのセクションは数学的な証明に割かれているので詳細はここで説明しないが、以下の結論が導かれている。

Th. 1

最初のゲームを行ったとき、解を求めるための最小のクエリ回数 $k$ は、以下の不等式を満たす最小の $k$ である。

$$ k^2 + k + 2 \leq 2^{k-m+1} $$

そしてこれに $m=6$ を当てはめて計算すると $k=13$ となるため、この問題は常に解くことができると分かる。しかし、論文には具体的な構成方法が記されていないため読み解く必要があった。

用語 (論文前半の説明)

このゲームには状態があるためそのための状態を $(a, b, c)$ と定義する。説明は以下の表にまとめる。

状態 説明
$a$ 現時点で手に入っている回答が全て真と仮定したときの、解の集合の大きさ
$b$ 現時点で手に入っている回答のうちちょうどひとつが偽と仮定したときの、解の集合の大きさ
$c$ 現時点で手に入っている回答のうちふたつが偽と仮定したときの、解の集合の大きさ

初期状態は $(2^m, 0, 0)$ となる。

ここで、ある状態 $(a, b, c)$ から得られる集合の部分集合をそれぞれ $(X, Y, Z)$ と置き、その濃度を $x, y, z$ とする。「 $X \cup Y \cup Z$ の中に求める解は存在しますか?」というクエリを $[x, y, z]$ queryと表記する。 状態 $(a, b, c)$ から発された $[x, y, z]$ query は、その結果次第で2つの状態に新たに遷移する。これらはそれぞれ以下のようになる。

$$ YES_{a, b, c}(x, y, z) = (x, a - x + y, b - y + z) $$

$$ NO_{a,b,c}(x,y,z) = (a-x, x+b-y, y+c-z) $$

$X, Y, Z$ は元の状態の部分集合なので、 $a \geq x, b \geq y, c \geq z $ が成り立つことに注意する。

$a$ について、yesが返ってくれば集合 $X$ に解が含まれているので、新しい状態の大きさは $x$ となる。noが返って来れば集合 $X$ の全要素は解ではないことが分かるので、新しい状態には $X$ の要素は除かれ、大きさは $a-x$ となる。

$b$ と $c$ については、嘘が存在する可能性を考慮する必要があることに注意する。

$b$ についてyesが返ってきたとき、「元々嘘を1回もついていなかったがこの回で嘘をついた」のであれば $a$ のnoの場合と等しいので新しいサイズは $a-x$ である。「すでに嘘を1回ついているがこの回では嘘をついていない」のであれば $Y$ の中に解が含まれているので新しい状態の大きさは $y$ となる。両方合わせて、 $a-x+y$ が新しい状態の大きさとなる*1

$b$ についてnoが返ってきたとき、「元々嘘を1回もついていなかったがこの回で嘘をついた」のであれば $a$ のyesの場合と等しいので新しいサイズは $x$ である。「すでに嘘を1回ついているがこの回では嘘をついていない」のであれば $Y$ の中に解が存在しないので新しい状態の大きさは $b-y$ となる。両方合わせて、 $x+b-y$ が新しい状態の大きさとなる。

$c$ については省略する。

また、別のアプローチとして 状態 $(a, b, c)$ は現時点までに嘘をついた数が異なるため、それを用いて重み付けを行うことができる。残り $j$ 回のクエリが残されているときの状態 $(a, b, c)$ に対する重み $w_{j}(a,b,c)$ を以下で定義する。

$$ \begin{align} w_{j}(a,b,c) &= a \left( \begin{pmatrix} j \\ 2 \end{pmatrix}+ \begin{pmatrix} j \\ 1 \end{pmatrix}+ \begin{pmatrix} j \\ 0 \end{pmatrix} \right) + b \left( \begin{pmatrix} j \\ 1 \end{pmatrix}+ \begin{pmatrix} j \\ 0 \end{pmatrix} \right) + c \begin{pmatrix} j \\ 0 \end{pmatrix} \\ &= a \dfrac{j^2+j+2}{2} + b(j+1) + c \end{align} $$

この重みを先ほどの状態遷移を合わせて記述すると、以下のようになる。

$$ w_{j}(a, b, c) = w_{j-1}(YES_{a,b,c}(x,y,z) + NO_{a,b,c}(x,y,z)) $$

また、状態 $(a, b, c)$ について関数 $ch$ を定義する。

$$ ch(a, b, c) = \mathrm{min} \lbrace k: w_{k}(a,b,c) \leq 2^k \rbrace $$

関数 $ch$ はゲームの勝敗を判定するための関数として使われる。

状態 $(a, b, c), ch(a, b, c) = j$ について、以下が成り立つ$[x, y, z]$ queryが存在する状態のことをbalancedと定義する。

$$ \left| w_{j-1}( YES_{a,b,c} (x, y, z) ) - w_{j-1}( NO_{a,b,c} (x, y, z) ) \right| \leq 1 $$

balanced は解の候補をほぼ半分にできるクエリが存在するか? ということを示している。

構成法

Th.2

自然数 $m \geq 3$ について、状態 $(1, m, {}_m \mathrm{C}_2 )$ から $ch(1, m, {}_m \mathrm{C}_2)$ 回以内のクエリで数を特定可能

Th.2 から Th.1 を示すことができるので、論文の半分は Th.2 の証明に割かれている。

クエリを「求める数の $j$ bit目は1か?」に限定する。すると、嘘を1回もついていない場合は1回の質問ごとに解の候補が半分になることになる。また、嘘を $i$ 回ついている場合は嘘を1回もついていない場合の解に、すでに判明しているbitのうち $i$ 個を選んで反転させたもの全ての集合が解の候補となる。したがって、 $i$ 回目のクエリの後の状態は $(2^{m-i}, i2^{m-i}, {}_i \mathrm{C}_{2}2^{m-i})$ となる。したがって、 $m$ 回クエリを飛ばすと状態はTh.2で示した状態になる。

以降、状態 $(1, m, {}_m \mathrm{C}_2 )$ からどのように特定するかのみ考える。

ここで重要になるのが balanced の概念。balanced は探索範囲を半分にするクエリのことなので、balanced を保つように変化させていけば良さそうだ。

$m=6$ のとき、状態 $(1, m, {}_m \mathrm{C}_2 ) = (1, 6, 15)$ から1回のクエリごとに評価値 $w_{j-1}(YES_{a,b,c}(x,y,z))$ と $w_{j-1}(NO_{a,b,c}(x,y,z))$ の差が最も少なくなるようなクエリの選び方をし続け、元の値を特定できるかを実験して確かめる。

実験コードは以下のものを使った。

def weight(a, k):
    return a[0] * (k**2 + k + 2) // 2 + a[1] * (k + 1) + a[2]

goal = ([1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0])
def calcquery(nowstate, zan):
    minquery = [-1, -1, -1]
    mincost = 10**8
    for i in range(nowstate[0] + 1):
        for j in range(nowstate[1] + 1):
            for k in range(nowstate[2] + 1):
                query = [i, j, k]
                yesstate = [query[0], nowstate[0] - query[0] + query[1], nowstate[1] - query[1] + query[2]]
                nostate = [nowstate[0] - query[0], query[0] + nowstate[1] - query[1], query[1] + nowstate[2] - query[2]]
                newcost = max(weight(yesstate, zan), weight(nostate, zan)) - min(weight(yesstate, zan), weight(nostate, zan))
                if mincost > newcost:
                    mincost = newcost
                    minquery = query
    return minquery

def jikken(nowstate, zan):
    if zan == 0:
        print(nowstate)
        if nowstate in goal:
            print('ok')
        else:
            print('ng')
        return
    query = calcquery(nowstate, zan)
    yesstate = [query[0], nowstate[0] - query[0] + query[1], nowstate[1] - query[1] + query[2]]
    nostate = [nowstate[0] - query[0], query[0] + nowstate[1] - query[1], query[1] + nowstate[2] - query[2]]
    jikken(yesstate, zan - 1)
    jikken(nostate, zan - 1)

jikken([1, 6, 15], 8)

実験を行った結果、クエリの返答が正常ならば必ず解は一意に定まることが分かった。これは $m=6$ の場合に限った検証だが、論文の証明を見る限り $m=6$ 以外の場合でも balanced な状態を常に保てるかが重要な事項であると考えられる。

一意に定まることが分かったので、後は上の戦略を実装すれば良い。しかし、実験では集合のサイズにのみ注目したのに対し、実装では集合の中身にまで木を配る必要があることに注意する。

  • 集合のクエリはありうる値を全て or で繋ぐことで得られる
  • クエリのサイズのみ合っていれば良いので、クエリを作るときは先頭から順番に選ぶだけで良い
  • クエリの残り回数が残っている時に解が求まってしまうと、例外処理しないとバグる

このあたりがハマりポイントだった。

実装したコードは以下の通りである。コメントなどを書いていないため非常に読みにくいのは申し訳ない。

from pwn import *
import string, hashlib, copy

# sock = remote('', )
sock = process(['python3', 'app.py'])
def proof_of_work(sock):
    sock.recvuntil(b'+')
    salt = sock.recv(16)
    print('salt:', salt)
    sock.recvuntil(b'= ')
    digest = sock.recvline().strip()
    sock.recvuntil(b': ')
    charlist = string.digits + string.ascii_letters
    l = len(charlist)
    for item in range(pow(l, 4)):
        buf = b''
        for _ in range(4):
            buf += charlist[item % l].encode()
            item = item // l
        if hashlib.sha256(buf + salt).hexdigest().encode() == digest:
            sock.sendline(buf)
            break
    return

# proof_of_work(sock)
sock.recvuntil(b'head?\n')

chest_id = [0, 1, 2, 3, 4, 5]

def initState(base):
    state = [[], [], []]
    state[0].append(copy.copy(base))
    for i in range(len(base)):
        tmp = copy.copy(base)
        tmp[i] = tmp[i] ^ 1
        state[1].append(tmp)
    for i in range(len(base)):
        for j in range(i+1, len(base)):
            tmp = copy.copy(base)
            tmp[i] = tmp[i] ^ 1
            tmp[j] = tmp[j] ^ 1
            state[2].append(tmp)
    return state

def next(nowstate, query, yes: bool):
    newstate = [[], [], []]
    diffset = [copy.deepcopy(nowstate[i]) for i in range(3)]
    for i in range(3):
        diffset[i] = diffset[i][len(query[i]):]
    if yes:
        newstate[0] = copy.deepcopy(query[0])
        newstate[1] = diffset[0] + query[1]
        newstate[2] = diffset[1] + query[2]
    else:
        newstate[0] = diffset[0]
        newstate[1] = query[0] + diffset[1]
        newstate[2] = query[1] + diffset[2]
    return newstate

def weight(a, k):
    return a[0] * (k**2 + k + 2) // 2 + a[1] * (k + 1) + a[2]

def calcQuery(nowstate, zan):
    nowstateSize = [len(item) for item in nowstate]
    minquery = [-1, -1, -1]
    mincost = 10**8
    for i in range(nowstateSize[0] + 1):
        for j in range(nowstateSize[1] + 1):
            for k in range(nowstateSize[2] + 1):
                query = [i, j, k]
                yesstate = [query[0], nowstateSize[0] - query[0] + query[1], nowstateSize[1] - query[1] + query[2]]
                nostate = [nowstateSize[0] - query[0], query[0] + nowstateSize[1] - query[1], query[1] + nowstateSize[2] - query[2]]
                newcost = max(weight(yesstate, zan), weight(nostate, zan)) - min(weight(yesstate, zan), weight(nostate, zan))
                if mincost > newcost:
                    mincost = newcost
                    minquery = query
    newquery = [[], [], []]
    for i in range(3):
        for j in range(minquery[i]):
            newquery[i].append(copy.copy(nowstate[i][j]))
    return newquery

def createQuery(c):
    candidates = c[0] + c[1] + c[2]
    if len(candidates) == 0:
        return '1'
    query = []
    for candidate in candidates:
        q = ''
        for i in range(len(candidate)):
            q += 'B'+str(i) + ' == ' + str(candidate[i])
            if i < len(candidate) - 1:
                q += ' and '
        query.append('( ' + q + ' )')
    return ' or '.join(query)

def round(sock):
    sock.recvuntil(b'twice!\n')
    baseresult = []
    for i in range(6):
        sock.recvuntil(b': ')
        query = 'B'+str(i)
        sock.sendline(query.encode())
        res = sock.recvline().strip()[:-1].split()[-1]
        baseresult.append(1 if eval(res.decode()) else 0)
    state = initState(baseresult)
    assert len(state[0]) == 1 and len(state[1]) == 6 and len(state[2]) == 15
    zan = 8
    while zan > 0:
        sock.recvuntil(b': ')
        query = calcQuery(state, zan)
        queryStr = createQuery(query)
        sock.sendline(queryStr.encode())
        res = sock.recvline().strip()[:-1].split()[-1]
        if res == b'words':
            print(sock.recvline())
            exit(1)
        if queryStr != '1':
            state = next(state, query, eval(res))
        zan -= 1
    sock.recvline()
    for i in range(3):
        if len(state[i]) > 0:
            sock.sendline(' '.join(list(map(str, state[i][0]))).encode())
            break
    line = sock.recvline()
    print(line)
    return

for ind in range(50):
    print(f'round {ind}')
    round(sock)

line = sock.recvline()
print(line)
sock.close()

後書き

このゲームに必勝法が存在するのは、このゲームが interactive であることに起因している。このゲームの non-interactive なバージョンは「2bitの誤り訂正能力を持つ符号を構成する問題」である。しかし論文中には、2bitの誤り訂正能力を持つ符号を任意の $ m $ について構成することは難しいとある。

また、これ自体は興味深い結果だが、評価関数 $abs(w_{j-1}( YES_{a,b,c} (x, y, z) ) - w_{j-1}( NO_{a,b,c} (x, y, z) ))$ の最小値を効率的に求めるアルゴリズムが存在しない。そのためこのアルゴリズムは低速で、競技プログラミング的な面白さはなさそうだと感じた。

最後に、元の論文はテキストの抽出が難しい組版で翻訳が面倒だったので、紫乃さん (@joisino_) が作成した論文翻訳支援サービス Readable を使用した。

readable.joisino.net

テキストとして抽出した部分だけをいい感じに翻訳するサービスなのかな〜と勝手に思って敬遠していたが、文字認識をガンガン使って翻訳していく想像よりずっとパワフルなソフトだった。組版がちゃんとしている論文ならコピペの方が信用できるというのは変わらないが、こういうコピペが通用しにくい論文を見開きで比較しながら翻訳できるというのは、とても良い体験だった *2 。今後も使っていければと思う。

*1:これだと解の範囲を超えそうだけど大丈夫だろうか

*2:書き込みも楽だし