Polictf 2017 - Tower.py

from pwn import *
from collections import defaultdict


def add_row(graph, y, row):

    for x, col in enumerate(row):
        true_x = x/2
        if col == " ": # Can do down or left
            if not x & 1:  # This is the left wall part
                # Can go left
                graph[(true_x, y)].add((true_x - 1, y))
                graph[(true_x - 1, y)].add((true_x, y))
            else: # This is the space
                # Can go down
                graph[(true_x, y)].add((true_x, y - 1))
                graph[(true_x, y - 1)].add((true_x, y))

    return graph


def parse_maze(maze):
    # Graph is an adjacency list: key(point):{point...}
    graph = defaultdict(set)
    # Flip the maze to make iteration go in order
    maze = maze[::-1][:-2]

    for i, row in enumerate(maze):
        # Chop off the extra bytes on the right
        row = row[:-2]
        graph = add_row(graph, i, row)

    return graph


def find_path(start, goal, graph, visited, path):
    if start == goal:
        return path, True

    neighbours = graph[start]
    for neighbour in neighbours:
        if neighbour not in visited:
            new_path = path[:]
            new_path.append(neighbour)
            new_visited = visited.copy()
            new_visited.add(neighbour)
            p, s = find_path(neighbour, goal, graph, new_visited, new_path)
            if s: return p, s

    return None, False


def convert_path(path):
    wasdpath = ""
    current = path[0]
    path = path[1:]
    for node in path:
        if node[0] > current[0]:
            wasdpath += "d"
        elif node[0] < current[0]:
            wasdpath += "a"
        elif node[1] > current[1]:
            wasdpath += "w"
        elif node[1] < current[1]:
            wasdpath += "s"
        else:
            raise Exception("WTF?")
        current = node
    return wasdpath


def solve():
    maze_raw = r.recvuntil("start")
    maze = maze_raw.split("\n")[:-1]
    start = (0, 0)
    r.recvuntil("goal: ")
    goal = tuple([int(x) for x in r.recvuntil("\n").split(",")])
    print maze_raw
    graph = parse_maze(maze)
    log.info("Goal: {}".format(goal))

    path, solved = find_path(start, goal, graph, {start}, [start])
    if solved:
        print "Solved!"
    else:
        print "WTF no solution"
        return

    wasdpath = convert_path(path)
    print wasdpath

    r.recv()
    r.sendline(wasdpath)

    r.interactive()


if __name__ == "__main__":
    host = "tower.chall.polictf.it"
    port = 31337
    r = remote(host, port)
    solve()