Auto-killing script for DEF CON 28 CTF challenge rhg

auto-killing.py

# -*- coding: utf-8 -*-
from pwn import *
from time import sleep
import requests
import json

global map_elems, token, myx, myy
TIMEOUT = 0.5
URL = 'http://10.13.37.1:8080/state.json'
# URL = 'http://10.168.4.66:8080/state.json'
my_items = {
    'USB_KEY': 0,
    'WIRELESS_CARD': 0,
    'APT_MALWARE': 0,
    'PHISHING_KIT': 0,
    'NDAY': 0,
    'ZERODAY': 0
}
ITEM_CODE_MAP = {
    'APT_MALWARE': 982365152,
    'USB_KEY': 2145013005,
    'PHISHING_KIT': 3023560243,
    'WIRELESS_CARD': 2002263477,
    'NDAY': 4025300686,
    'ZERODAY': 2775105017
}
ATT_TOOLS = ['USB_KEY', 'WIRELESS_CARD', 'APT_MALWARE', 'PHISHING_KIT']
DAY_TOOLS = ['NDAY', 'ZERODAY']
L = 4271926414
R = 4265964054
U = 1026868169
D = 3930400191

direction = [
    [1, 0],
    [0, 1],
    [-1, 0],
    [0, -1]
]
char = 'SDWA'

to_x = 0
to_y = 0
cnt = 0
path = []
mp = []
dis = []
parent = []

def get_token(team_id):
    FLAG = 1
    import requests
    if FLAG:
        token = requests.get("http://123.206.180.189:9999/get_token/{}".format(team_id)).text
    else:
        token = sys.argv[1]
    return token

def submit_flag(flag):
    ret = requests.post('http://10.13.37.1/api/submit_flag/' + flag)
    return ret.text

def get_map_from_server():
    req = requests.get(URL)
    if req.status_code == 200:
        res = req.json()
        map_elems = res['elems']
        # print res
        return map_elems
    return None

def send_cmd(cmd_id, content=None):
    if content:
        p.sendline('{} {} {}'.format(cmd_id, token, str(content)))
    else:
        p.sendline('{} {}'.format(cmd_id, token))
    sleep(TIMEOUT)
    r = p.recvline()
    if 'rate limit' in r:
        log.info('rate limit')
        r = send_cmd(cmd_id, content)
    elif '"ERROR"' in r:
        log.error(r)
        pause()
        exit(0)
    return r

def check_flag(r):
    # 检测捡到的是不是flag
    resp = json.loads(r)
    if resp['status'] == 'OK':
        item = resp['info']['item']
        if item['item_type'] == 'FLAG':
            r = send_cmd(1083909441, item['id'])    # inspect flag
            insp_r = json.loads(r)
            if insp_r['status'] == 'OK':
                flag = insp_r['info']
                log.success(flag)
                print submit_flag(flag)
        else:
            my_items[item['item_type']] += 1

def do_pick(d=None):
    '''
    捡上下左右或者一个指定方向
    '''
    r = send_cmd(1169593071, d)
    print r
    check_flag(r)

def do_attack(d=None):
    for i in ATT_TOOLS:
        if my_items[i] > 0:
            if my_items['NDAY'] > 0:
                cmd = '{} {} '.format(ITEM_CODE_MAP[i], ITEM_CODE_MAP['NDAY'])
                break
            elif my_items['ZERODAY'] > 0:
                cmd = '{} {} '.format(ITEM_CODE_MAP[i], ITEM_CODE_MAP['ZERODAY'])
                break
            else:
                print 'not enough item, can\'t attack'
                return
    else:
        print 'not enough item, can\'t attack'
        return
    print cmd + str(d)
    r = send_cmd(2257090568, cmd + str(d))
    print r
    # do_pick(str(d))
    # do_pick(U)
    # do_pick(D)
    # do_pick(L)
    # do_pick(R)

def up():
    global myx, myy
    print 'up'
    r = send_cmd(2614795397, U)
    print r
    if '"status": "OK"' in r:
        myy -= 1

def down():
    global myx, myy
    print 'down'
    r = send_cmd(2614795397, D)
    print r
    if '"status": "OK"' in r:
        myy += 1

def left():
    global myx, myy
    print 'left'
    r = send_cmd(2614795397, L)
    print r
    if '"status": "OK"' in r:
        myx -= 1

def right():
    global myx, myy
    print 'right'
    r = send_cmd(2614795397, R)
    print r
    if '"status": "OK"' in r:
        myx += 1

def get_item_loc(l):
    global myx, myy, map_elems
    min_dis = 0xffff
    for k, v in map_elems.items():
        if v.get('item_type', None) in l:
            if v.get('loc', None) is None:
                continue
            tmpy, tmpx = v['loc']
            dis = ((myx - tmpx) ** 2) + ((myy - tmpy) ** 2)
            if dis < min_dis:
                min_dis = dis
                target_y, target_x = tmpy, tmpx
    return target_y, target_x

def bfs(x, y):
    mp[to_x][to_y] = ' '
    q = []
    q.append([x,y])
    while(len(q) != 0):
        x, y = q[0]
        q.pop(0)
        if x == to_x and y == to_y:
            return True
        for i in range(0, 4):
            new_x = x + direction[i][0]
            new_y = y + direction[i][1]
            if new_x > 29 or new_x < 0 or new_y > 29 or new_y < 0 or \
                mp[new_x][new_y] != ' ':
                continue
            q.append([new_x, new_y])
            dis[new_x][new_y] = dis[x][y] + 1
            mp[new_x][new_y] = char[i]
            parent[new_x][new_y] = [x, y]
    return False

def find_path(x, y, to_x, to_y):
    global path
    while to_x != x or to_y != y:
        path = [mp[to_x][to_y]] + path
        to_x, to_y = parent[to_x][to_y]

def deal(map_elems, fr, to):
    global mp
    global dis
    global parent
    global to_x
    global to_y
    global path
    path = []
    mp = []
    dis = []
    parent = []
    for i in range(0, 30):
        mp.append([])
        dis.append([])
        parent.append([])
        for j in range(0, 30):
            mp[i].append(' ')
            dis[i].append(0x7f7f7f7f)
            parent[i].append(0)
    for key in map_elems:
        if map_elems[key]['loc'] and len(map_elems[key]['loc']) != 0:
            x, y = map_elems[key]['loc']
            mp[x][y] = '.'
    x, y = fr
    to_x, to_y = to
    bfs(x, y)
    find_path(x,  y, to_x, to_y)
    return path

def go_to_and_do(myx, myy, target_x, target_y, callback):
    print 'now my location', myx, myy
    print 'target location', target_x, target_y
    path = deal(map_elems, [myy, myx], [target_y, target_x])
    print path, 'path len:', len(path)
    for i, d in enumerate(path):
        print 'Moving', i + 1
        if d == 'A':
            if i == len(path) - 1:
                print 'callback to dir: left'
                callback(L)
            else:
                left()
        elif d == 'D':
            if i == len(path) - 1:
                print 'callback to dir: right'
                callback(R)
            else:
                right()
        elif d == 'W':
            if i == len(path) - 1:
                print 'callback to dir: up'
                callback(U)
            else:
                up()
        elif d == 'S':
            if i == len(path) - 1:
                print 'callback to dir: down'
                callback(D)
            else:
                down()

def get_my_items():
    r = send_cmd(3623901639)
    items = json.loads(r)
    for k, v in items['info']['items'].items():
        if v.get('item_type', None) in ATT_TOOLS + DAY_TOOLS:
            my_items[v['item_type']] += 1

def update_map():
    global map_elems
    r = send_cmd(1415591046)
    map_elems.update(json.loads(r)['info']['elems'])


src_team = 9
dst_team = 10
if len(sys.argv) > 1:
    dst_team = int(sys.argv[1])
is_remote = False
is_remote = True
token = get_token(src_team)

# context.log_level = 'debug'
if is_remote:
    p = remote('10.13.37.1', 14000)
    p.sendlineafter('auth token> ', '1f6a1f6a1f6a1f6a1f6a1f6a1f6a1f6a')
    # p.sendline('cat /var/rhg/token')
    # token = p.recv(0.50)
    # print token
    # token = sys.argv[1]
    p.sendline('nc 10.0.70.100 8000')
else:
    p = remote('10.168.4.66', 6666)
print p.recvline()
p.sendline('AUTH ' + token + ' ' + token)
sleep(TIMEOUT)
print p.recvline()

map_elems = get_map_from_server()
get_my_items()
myy, myx = map_elems['P' + str(src_team)]['loc']

if all([my_items[i] == 0 for i in ATT_TOOLS]):
    target_y, target_x = get_item_loc(ATT_TOOLS)
    go_to_and_do(myx, myy, target_x, target_y, do_pick)

if all([my_items[i] == 0 for i in ['NDAY']]):
    target_y, target_x = get_item_loc(['NDAY'])
    go_to_and_do(myx, myy, target_x, target_y, do_pick)

# attack enemy
enemyy, enemyx = map_elems['P' + str(dst_team)]['loc']
go_to_and_do(myx, myy, enemyx, enemyy, do_attack)

update_map()
print 'updating map to get flag location'
target_y, target_x = get_item_loc(['FLAG'])
go_to_and_do(myx, myy, target_x, target_y, do_pick)