kompreplicants tower of hanoi

decode.py

# the encoding is
# first four bits are the depth in the tree encoded
# next 8 bits is the character encoded
# next X bits is the position in the tree encoded with the depth given by the first four bits


# The file has the following structure
# Each tree node encoded + 4 bits set to zero + each original character encoded by its position in the tree + few bits at the end

bits = []
byts = []

# Get the bytes in the file
with open("flag.txt.short008", "rb") as f:
   
    b = f.read(1)
    while b != '':
        byts.append(ord(b))
        b = f.read(1)


# Get the bits from the bytes and re-order them to get them in the order they were emmitted to the writing routine
for i, b in enumerate(byts):
    #print bin(b), hex(b)
    tmp = [0] * 8
    s = ""
    if i == 0:
        s = bin(b)[2:]
    else:
        if b < byts[i-1]:
            s = bin(b + (256 - byts[i - 1]))[2:]
        else:
            s = bin(b - byts[i - 1])[2:]
    for i, c in enumerate(s):
        tmp[8 - len(s) +i] = int(c)
    #print tmp
    bits += tmp



#print bits

#global index into the bit stream
i = 0

# read the depth by reading 4 bits and rebuilding the number
def get_depth():
    global i    
    depth = bits[i + 3] + (bits[i+2] << 1) + (bits[i+1] << 2) + (bits[i] << 3)
    i += 4
    return depth


# get the encoded character by reading 8 bits
def get_char():
    global i
    char = bits[i] << 7
    char += bits[i+1] << 6
    char += bits[i+2] << 5
    char += bits[i+3] << 4
    char += bits[i+4] << 3
    char += bits[i+5] << 2
    char += bits[i+6] << 1 
    char += bits[i+7]
    #print hex(char)
    i += 8
    return char

dct2 = {}
dct = {}
len_pos = 1

# read a node from the tree i.e get the depth, the character, and its position
def read_tree_node():
    global i
    global dct
    global dct2
    global len_pos
    d = get_depth()
    c = get_char()
    #print 'depth: %d' % d
    pos = 0
    for j in xrange(0, d):
        pos += bits[i] << j
        i += 1
    if c != 0x80:
        dct[pos] = c
        dct2[pos] = d


# lazy way, after writing out the tree representation
# the binary writes four 0 bits therefore any sequence of four bits
# set to zero is potentially the limit between the tree and the encoded chars
# get all the possible ends and we'll parse the file assuming the end is at the index being iterated
def find_potential_end():
    l = []
    k = 0 
    while k < len(bits) - 5:
        if bits[k] == 0 and bits[k+1] == 0 and bits[k+2] == 0 and bits[k+3] == 0:
            #print 'found potential end at indexi %d' % (k)
            l.append(k)
        k += 1
    return l


# this is a really hackish function
# in the midst of everything i did not think of any simpler way
# this is the result of a lot of broken implementation
# (read that as bad code on tope of bad code on top of bad code)
# when we read the encoded char (i.e. its position in the tree)
# we'll check if there is only one single position that has the same
# binary representation
# for example, lets say the position are the following
# 00 11 101
# then if we read 0 we know that the position is going to be 00 
# and we need to advance in the bit stream as such (done by the calling function)
# if we read 1 though, the position can be 11 or 101 so we read a next byte etc..
def is_single_occurence(number, length):
    cnt = 0
    res = -1
    #print 'called with', number
    for key in dct:
        
        s1 = bin(key)[2:][::-1].ljust(length, '0')
        s2 = bin(number)[2:][::-1].ljust(length, '0')
        #print s1, s2
        if s1.startswith(s2):
            cnt += 1
            res = key
            if cnt == 2:
                return (False, None)
        
    if cnt == 1:
        return (True, res)
    return (False, None)
    

# called after rebuilding the tree, tries to rebuild the flag by reading each encoded character
def rebuild_flag(i):
    flag = ""
    while i < len(bits):
        try:
            shift = 0
            pos = 0
            pos += bits[i]
            i += 1
            while not is_single_occurence(pos, shift + 1)[0]:
                shift += 1
                pos += bits[i] << shift
                i += 1
            for j in xrange(shift + 1, dct2[is_single_occurence(pos, shift + 1)[1]]):
                shift += 1
                pos += bits[i] << shift
                i += 1
            flag += chr(dct[pos])
        except Exception as e:
            print e
            print 'FLAG', flag
            return
    print 'FLAG', flag


def rebuild_tree(end):
    global i
    char = 0
    while i < end:
        read_tree_node()
    flag = ''
    i = end + 4
    rebuild_flag(i)

#rebuild_tree()
for end in find_potential_end():
    print 'trying end at index %d...' % (end)
    i = 0
    rebuild_tree(end)
    dct = {}
    dct2 = {}

#python decode.py > out && grep flag{ out
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}
#flag{"m0r3_c0mpr3553d_7h4n_6unz1p" 15 0ur m0770}