main.py

import asyncio
import uuid

import uvicorn
from typing import Dict, Optional

from fastapi import FastAPI, APIRouter, UploadFile, BackgroundTasks
from fastapi.templating import Jinja2Templates
from loguru import logger
from starlette.requests import Request
from starlette.responses import HTMLResponse
from starlette.staticfiles import StaticFiles

from utils import parse_line
from core import AsyncGrassWs

app = FastAPI()

templates = Jinja2Templates(directory="app/templates")
app.mount("/static", StaticFiles(directory="app/static"), name="static")

client_router = APIRouter(prefix='/client')

all_client: Dict[str, AsyncGrassWs] = {}
all_client_ids = []

CLIENT_INDEX = 0

# Or if there are multiple tasks
background_tasks = set()


def run_client(client_id):
    task = asyncio.create_task(all_client[client_id].run())

    # Add the task to the set to keep a strong reference:
    background_tasks.add(task)

    # To prevent the task from being held in memory forever and unable to be garbage collected
    # Let each task remove itself from the set after it finishes:
    task.add_done_callback(background_tasks.discard)
    return client_id


def add_client(grass_client: AsyncGrassWs):
    client_id = uuid.uuid4().__str__()
    all_client[client_id] = grass_client
    all_client_ids.append(client_id)
    return client_id


async def delete_client(client_id):
    logger.info(f'[Exit] {all_client[client_id].user_id}')
    await all_client[client_id].stop()
    del all_client[client_id]
    all_client_ids.remove(client_id)


def load_file_clients(data):
    new_clients = []
    index = 0
    for line in data.split('\n'):
        user_id, proxy_url = parse_line(line)
        if not user_id:
            continue
        index += 1
        client = AsyncGrassWs(user_id=user_id, proxy_url=proxy_url)
        new_clients.append(add_client(client))
    return new_clients


async def threading_run_clients(clients):
    for client_id in clients:
        run_client(client_id)


@app.get("/", response_class=HTMLResponse)
async def read_item(request: Request):
    return templates.TemplateResponse(request=request, name="index.html")


@client_router.get("/{client_id}")
def find_one(client_id: str):
    client = all_client.get(client_id)
    data = {
        'data': {
            'status': None,
            "proxy_url": None,
            "logs": []
        },
        'message': "failed"
    }
    if client is not None:
        data = {
            'data': {
                'status': client.status,
                "proxy_url": client.proxy_url,
                "logs": list(reversed(client.logs[-50:]))
            },
            'message': "success"
        }
    return data


@client_router.get("/")
def find_all():
    data = []
    for client_id in all_client_ids:
        try:
            data.append({
                'id': client_id,
                'user_id': all_client[client_id].user_id,
                'status': all_client[client_id].status,
                "proxy_url": all_client[client_id].proxy_url
            })
        except:
            continue
    return {
        'data': data,
        'message': "success"
    }


@client_router.post("/")
async def add(user_id: str, proxy_url: Optional[str] = None):
    client = AsyncGrassWs(user_id=user_id, proxy_url=proxy_url or None)
    client_id = add_client(client)
    run_client(client_id)
    return {'data': client_id, 'message': 'create success'}


@client_router.delete("/{user_id}")
async def delete_one(user_id: str):
    await delete_client(user_id)
    return {'data': user_id, 'message': 'success'}


@client_router.delete("/")
async def delete_all():
    all_client_ids_copy = all_client_ids[::]
    for client_id in all_client_ids_copy:
        await delete_client(client_id)
    return {'data': [], 'message': 'success'}


@app.post("/upload/")
async def run_by_file(file: UploadFile, background_task: BackgroundTasks):
    data = (await file.read()).decode()
    new_clients = load_file_clients(data)
    background_task.add_task(threading_run_clients, new_clients)
    return {"data": None, 'message': 'success'}


app.include_router(client_router)

if __name__ == '__main__':
    uvicorn.run("main:app", host='0.0.0.0')