Stegasaurus Ccratch solution (PlaidCTF 2020)

part1.py

# The solution comes from the paper https://sci-hub.tw/10.1007/BF03025305 
# Which I got from p4 team.

import random
from math import factorial

SET_SIZE = 8
MAX_VAL = 40000

# get random 8 integers
def get_random_8():
    ret = []
    while len(ret) < 8:
        x = random.randint(0, MAX_VAL)
        if x not in ret:
            ret.append(x)
    return ret

# convert permutation to an integer
def perm_to_int(perm):
    perm = perm[::-1]
    N = len(perm)
    a = sorted(perm)
    p = [a.index(x) for x in perm]
    ret = 0
    for i in range(N):
        ret += factorial(p[i]) * len( [None for x in p[:i] if x < p[i]] )

    return ret

# convert an integer to permutation of elements from arr
def int_to_perm(v, arr):
    N = len(arr)
    assert v < factorial(N)
    tmp = []
    for x in range(N-1, -1, -1):
        tmp += [ v // factorial(x) ]
        v = v % factorial(x)

    tmp = tmp[::-1]
    ret = []
    for i in range(N):
        p = tmp[i]
        ret = ret[:p] + [i] + ret[p:]

    return list(map(lambda x: arr[x], ret[::-1]))

# let the discarded element be the element of the index idx which is calculated
# by idx = sum(numbers_8) mod 8. Notice that then, the number is congurent to 
# idx - sum(numbers_7) mod 8, where numbers_7 is the array without discarded elem
# Alice can send to bob discarded // 8, which is less than 5040 (7!) and which can be 
# encoded as the permutation of numbers. 

def alice(arr):
    N = len(arr)
    a = sorted(arr)
    discarded = a[sum(a) % N]
    a.remove(discarded)
    mod = discarded // SET_SIZE

    perm = int_to_perm(mod, a)

    return perm, discarded

# Bob recovers the original number which is close to the discarded one. 
# The difference is the reminder discarded % 8. 
# Because Alice encoded the reminder by choosing which element to disacrd
# by calculating idx, Bob can find how many elements were less than that number
# and therefore find idx mod 8, which is the last missing piece.

def bob(arr):
    val = perm_to_int(arr)

    for i in range(SET_SIZE):
        num_candidate = SET_SIZE*val + (-sum(arr) + i)% SET_SIZE
        if sum(x <= num_candidate for x in arr) == i:
            return num_candidate
    

fails = 0

# validate the correctness of the algorithm
for _ in range(10000):
    arr, discarded = alice(get_random_8())
    if discarded != bob(arr):
        fails += 1

print("Fails: %d" % fails)

part2.py

# The idea for the solution is the observation that at least one 2 must be adjacent to one 1. 
# So the Alice wants to choose that 1 because Bob can then know that there must be 2 next to it.
# By induction, that one pair of (1,2) can be "removed" and the process repeated, which proves
# solvability of the problem. If we treat the array as a cycle, then there are exactly two 2 that
# are adjacent to 1. I chose to always leave the 1 that is right to 2, e.g. 22221

import random

# for bob, we want to always put 2 next to 1 from left, (e.g. 021)
# if we don't have any free spot left, then we are looping until we find the free spot
# treating the array as a cycle. 

def get_random_sol():
    x = [1]*64 + [2]*32
    random.shuffle(x)
    return x

def bob(a):
    arr = a[::]
    N = len(arr)
    for i in range(N):
        # if the filled is not 1, we skip it
        if arr[i] != 1:
            continue
        j = i

        # find free spot
        while arr[j] != 0:
            j = (j-1) % N
        
        arr[j] = 2
    arr = list(map(lambda x: 1 if x == 0 else x, arr))
    return arr


# for alice we want to have reverse function to bob
# now we search for 2, and put 1 right to it

def alice(a):
    arr = list(map(lambda x: 0 if x == 1 else 2, a))

    N = len(arr)
    for i in range(N):
        if arr[i] != 2:
            continue
        j = i
        while arr[j] != 0:
            j = (j+1) % N
        arr[j] = 1

    arr = list(map(lambda x: 0 if x == 2 else x, arr))
    assert a == bob(arr), "Nope"
    return arr


for _ in range(10000):
    x = get_random_sol()
    assert x == bob(alice(x)), "smth wrong"