From 369802ecb6ef0f87218b8eefd9a21979f3494f93 Mon Sep 17 00:00:00 2001 From: mradigen <55953083+mradigen@users.noreply.github.com> Date: Tue, 14 Nov 2023 09:31:53 +0530 Subject: [PATCH] chore: Fix config.py --- src/pwncore/__init__.py | 3 +- src/pwncore/config.py | 81 ++++++++++++++++----------------- src/pwncore/routes/ctf/start.py | 24 +++++----- 3 files changed, 52 insertions(+), 56 deletions(-) diff --git a/src/pwncore/__init__.py b/src/pwncore/__init__.py index 2b47838..867c3dd 100644 --- a/src/pwncore/__init__.py +++ b/src/pwncore/__init__.py @@ -4,6 +4,7 @@ import pwncore.docs as docs import pwncore.routes as routes +from pwncore.config import config app = FastAPI( title="Pwncore", openapi_tags=docs.tags_metadata, description=docs.description @@ -13,7 +14,7 @@ register_tortoise( app, - db_url="sqlite://:memory:", + db_url=config.db_url, modules={"models": ["pwncore.db"]}, generate_schemas=True, add_exception_handlers=True diff --git a/src/pwncore/config.py b/src/pwncore/config.py index 85d0ee6..48040bd 100644 --- a/src/pwncore/config.py +++ b/src/pwncore/config.py @@ -1,43 +1,38 @@ -from __future__ import annotations - -import typing as t - -__all__ = ( - "Config", - "BaseConfig", -) - - -class Config(t.Protocol): - development: bool - flag: str - max_containers_per_team: int - messages: dict - - -class BaseConfig(Config): - __slots__ = ("development",) - - flag = "C0D" - max_containers_per_team = 3 - messages = { - "db_error": "An error occured, please try again.", - - "port_limit_reached": "Server ran out of ports 💀", - - "ctf_not_found": "CTF does not exist.", - - "container_start": "Container started.", - "container_stop": "Container stopped.", - "containers_team_stop": "All team containers stopped.", - "container_not_found": "You have no running containers for this CTF.", - "container_already_running": "Your team already has a running container for this CTF.", - "container_limit_reached": "Your team already has reached the maximum number of containers limit, please stop other unused containers." - } - - def __init__(self, development: bool) -> None: - self.development = development - -DEV_CONFIG: t.Final[BaseConfig] = BaseConfig(True) - -config = DEV_CONFIG +from dataclasses import dataclass + +""" +Sample messages: +"db_error": "An error occurred, please try again.", +"port_limit_reached": "Server ran out of ports 💀", +"ctf_not_found": "CTF does not exist.", +"container_start": "Container started.", +"container_stop": "Container stopped.", +"containers_team_stop": "All team containers stopped.", +"container_not_found": "You have no running containers for this CTF.", +"container_already_running": "Your team already has a running container for this CTF.", +"container_limit_reached": "Your team already has reached the maximum number of containers limit, please stop other unused containers." +""" + +msg_codes = { + "db_error": 0, + "port_limit_reached": 1, + "ctf_not_found": 2, + "container_start": 3, + "container_stop": 4, + "containers_team_stop": 5, + "container_not_found": 6, + "container_already_running": 7, + "container_limit_reached": 8 +} + + +@dataclass +class Config: + development: bool + msg_codes: dict + db_url: str = "sqlite://:memory:" + flag: str = "C0D" + max_containers_per_team: int = 3 + + +config = Config(development=True, msg_codes=msg_codes) diff --git a/src/pwncore/routes/ctf/start.py b/src/pwncore/routes/ctf/start.py index 6d62bd6..a8e0603 100644 --- a/src/pwncore/routes/ctf/start.py +++ b/src/pwncore/routes/ctf/start.py @@ -35,20 +35,20 @@ async def start_docker_container(ctf_id: int, response: Response): ctf = await CTF.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 - return {"msg": config.messages["ctf_not_found"]} + return {"msg_code": config.msg_codes["ctf_not_found"]} team_id = get_team_id() # From JWT team_container = await Container.get_or_none(team_id=team_id, ctf_id=ctf_id) if team_container: return { - "msg": config.messages["container_already_running"], + "msg_code": config.msg_codes["container_already_running"], "ports": team_container.ports.split(","), "ctf_id": team_container.ctf_id } if await Container.filter(team_id=team_id).count() >= config.max_containers_per_team: return { - "msg": config.messages["container_limit_reached"] + "msg_code": config.msg_codes["container_limit_reached"] } # Start a new container @@ -61,7 +61,7 @@ async def start_docker_container(ctf_id: int, response: Response): # Handle error here logging.critical("No more free ports available on machine.") response.status_code = 500 - return {"msg": config.messages["port_limit_reached"]} + return {"msg_code": config.msg_codes["port_limit_reached"]} ports = [] # Only to save the host ports used to return to the user for guest_port in image_config["ports"]: @@ -99,11 +99,11 @@ async def start_docker_container(ctf_id: int, response: Response): response.status_code = 500 return { - "msg": config.messages["db_error"] + "msg_code": config.msg_codes["db_error"] } return { - "msg": config.messages["container_start"], + "msg_code": config.msg_codes["container_start"], "ports": ports, "ctf_id": ctf_id } @@ -124,7 +124,7 @@ async def stopall_docker_container(response: Response): except Exception: response.status_code = 500 return { - "msg": config.messages["db_error"] + "msg_code": config.msg_codes["db_error"] } for db_container in containers: @@ -132,7 +132,7 @@ async def stopall_docker_container(response: Response): container.stop() container.remove() - return {"msg": config.messages["containers_team_stop"]} + return {"msg_code": config.msg_codes["containers_team_stop"]} @atomic() @@ -142,12 +142,12 @@ async def stop_docker_container(ctf_id: int, response: Response): ctf = await CTF.get_or_none(id=ctf_id) if not ctf: response.status_code = 404 - return {"msg": config.messages["ctf_not_found"]} + return {"msg_code": config.msg_codes["ctf_not_found"]} team_id = get_team_id() team_container = await Container.get_or_none(team_id=team_id, ctf_id=ctf_id) if not team_container: - return {"msg": config.messages["container_not_found"]} + return {"msg_code": config.msg_codes["container_not_found"]} # We first try to delete the record from the DB # Then we stop the container @@ -156,11 +156,11 @@ async def stop_docker_container(ctf_id: int, response: Response): except Exception: response.status_code = 500 return { - "msg": config.messages["db_error"] + "msg_code": config.msg_codes["db_error"] } container = docker_client.containers.get(team_container.id) container.stop() container.remove() - return {"msg": config.messages["container_stop"]} + return {"msg_code": config.msg_codes["container_stop"]}