#!/usr/bin/env python3
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from aiohttp import web
import capstone
import functools
from gdbproc import GDBProcess
import socketio
import asyncio
import codecs
import os

enable_logging = False

premium = 'PREMIUM' in os.environ
if premium:
  access_key = os.getenv('PREMIUM_KEY')
  runnable = ['/home/user/printwebflag']
else:
  access_key = os.getenv('TRIAL_KEY')
  runnable = ['/bin/sleep', '20']

MAX_INSN_LEN = 15

capstone_md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)

sio = socketio.AsyncServer()
app = web.Application()
sio.attach(app)

with open('index.html') as f:
  index_html = f.read()

async def index(request):
  if not 'key' in request.cookies:
    return web.Response(status=401, text='permission denied (missing key)', content_type='text/html')
  if request.cookies['key'] != access_key:
    return web.Response(status=401, text='permission denied (invalid key)', content_type='text/html')
  return web.Response(text=index_html, content_type='text/html')

app.add_routes([web.get('/', index),
                web.get('/{name}', index)])

gdb_sessions = {}
stop_queue_readers = {}

async def on_shutdown(app):
  await asyncio.gather(delete_gdb_process(sid) for sid in gdb_sessions.keys())

app.on_shutdown.append(on_shutdown)

def log(msg):
  if enable_logging:
    print('[*] {}'.format(msg))

@sio.on('connect')
def connect(sid, environ):
  log('connected {}'.format(sid))
  if not 'key={}'.format(access_key) in environ['HTTP_COOKIE']:
    log('access_key not found {}'.format(environ['HTTP_COOKIE']))
    return False

@sio.on('disconnect')
async def disconnect(sid):
  log('disconnected {}'.format(sid))
  await delete_gdb_process(sid)

async def stop_queue_reader(sid, queue):
  while True:
    pkt = await queue.get()
    await update_all(sid)

async def create_gdb_process(sid):
  stop_queue = asyncio.Queue()
  gdb_sessions[sid] = await GDBProcess.create(runnable, stop_queue, env={'KEY': access_key}, log_fn=log)
  loop = asyncio.get_event_loop()
  stop_queue_readers[sid] = loop.create_task(stop_queue_reader(sid, stop_queue))

async def delete_gdb_process(sid):
  if sid in gdb_sessions:
    stop_queue_readers[sid].cancel()
    del stop_queue_readers[sid]
    await gdb_sessions[sid].release()
    del gdb_sessions[sid]

@sio.on('start')
async def start(sid):
  await delete_gdb_process(sid)
  await create_gdb_process(sid)
  # Reading registers doesn't work on ubuntu 18.04 for some reason.
  # Step once as a work around
  step(sid)

async def update_all(sid):
  log('updating sid {}'.format(sid))
  regs_task = getregs(sid)
  maps_task = getmaps(sid)
  asm_task = getasm(sid, {'addr': await gdb_sessions[sid].get_reg('rip'), 'count': 100})
  await asyncio.gather(regs_task, maps_task, asm_task)
  log('update done')

@sio.on('step')
def step(sid):
  gdb_sessions[sid].step()

@sio.on('cont')
def cont(sid):
  gdb_sessions[sid].cont()

@sio.on('stop')
def stop(sid):
  gdb_sessions[sid].interrupt()

async def getregs(sid):
  regs = await gdb_sessions[sid].get_regs()
  await sio.emit('regs', regs, room=sid)

@sio.on('mem')
async def getmem(sid, msg):
  addr = msg['addr']
  count = msg['count']
  data = gdb_sessions[sid].read_mem(addr, count)
  await sio.emit('mem', {'addr': addr, 'data': data}, room=sid)

async def getmaps(sid):
  maps = gdb_sessions[sid].maps()
  await sio.emit('maps', maps, room=sid)

@sio.on('break')
async def setbreakpoint(sid, data):
  addr = data['addr']
  await gdb_sessions[sid].set_breakpoint(addr)
  await sio.emit('breakpoints', gdb_sessions[sid].breakpoints(), room=sid)

@sio.on('unbreak')
async def rmbreakpoint(sid, data):
  addr = data['addr']
  await gdb_sessions[sid].remove_breakpoint(addr)
  await sio.emit('breakpoints', gdb_sessions[sid].breakpoints(), room=sid)

@sio.on('search')
async def search(sid, data):
  q = data['q']
  qtype = data['type']
  await sio.emit('search_result', gdb_sessions[sid].search(q.encode(), qtype), room=sid)

async def getasm(sid, data):
  addr = data['addr']
  count = data['count']
  result = []
  for _ in range(count):
    data = gdb_sessions[sid].read_mem(addr, MAX_INSN_LEN)
    try:
      disasm = next(capstone_md.disasm_lite(data, addr))
    except StopIteration:
      break
    result.append(disasm)
    addr += disasm[1]
  await sio.emit('asm', result, room=sid)

if __name__ == '__main__':
  web.run_app(app)

[((38, 14, 38, 68), 'capstone.Cs', 'capstone.Cs', ({(38, 26, 38, 46): 'capstone.CS_ARCH_X86', (38, 48, 38, 67): 'capstone.CS_MODE_64'}, {}), '(capstone.CS_ARCH_X86, capstone.CS_MODE_64)', False, 'import capstone\n'), ((40, 6, 40, 28), 'socketio.AsyncServer', 'socketio.AsyncServer', ({}, {}), '()', False, 'import socketio\n'), ((41, 6, 41, 23), 'aiohttp.web.Application', 'web.Application', ({}, {}), '()', False, 'from aiohttp import web\n'), ((30, 15, 30, 39), 'os.getenv', 'os.getenv', ({(30, 25, 30, 38): '"""PREMIUM_KEY"""'}, {}), "('PREMIUM_KEY')", False, 'import os\n'), ((33, 15, 33, 37), 'os.getenv', 'os.getenv', ({(33, 25, 33, 36): '"""TRIAL_KEY"""'}, {}), "('TRIAL_KEY')", False, 'import os\n'), ((52, 9, 52, 64), 'aiohttp.web.Response', 'web.Response', (), '', False, 'from aiohttp import web\n'), ((87, 15, 87, 30), 'asyncio.Queue', 'asyncio.Queue', ({}, {}), '()', False, 'import asyncio\n'), ((89, 9, 89, 33), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((175, 2, 175, 18), 'aiohttp.web.run_app', 'web.run_app', ({(175, 14, 175, 17): 'app'}, {}), '(app)', False, 'from aiohttp import web\n'), ((49, 11, 49, 101), 'aiohttp.web.Response', 'web.Response', (), '', False, 'from aiohttp import web\n'), ((51, 11, 51, 101), 'aiohttp.web.Response', 'web.Response', (), '', False, 'from aiohttp import web\n'), ((54, 16, 54, 35), 'aiohttp.web.get', 'web.get', ({(54, 24, 54, 27): '"""/"""', (54, 29, 54, 34): 'index'}, {}), "('/', index)", False, 'from aiohttp import web\n'), ((55, 16, 55, 41), 'aiohttp.web.get', 'web.get', ({(55, 24, 55, 33): '"""/{name}"""', (55, 35, 55, 40): 'index'}, {}), "('/{name}', index)", False, 'from aiohttp import web\n'), ((88, 28, 88, 104), 'gdbproc.GDBProcess.create', 'GDBProcess.create', (), '', False, 'from gdbproc import GDBProcess\n'), ((112, 8, 112, 54), 'asyncio.gather', 'asyncio.gather', ({(112, 23, 112, 32): 'regs_task', (112, 34, 112, 43): 'maps_task', (112, 45, 112, 53): 'asm_task'}, {}), '(regs_task, maps_task, asm_task)', False, 'import asyncio\n')]