diff --git a/.gitignore b/.gitignore index f2c8323..37de172 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ __pycache__/ *.egg-info/ build/ dist/ +*.yml +*.yaml diff --git a/Dockerfile b/Dockerfile index bb2755a..8890311 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,4 +3,4 @@ WORKDIR /app COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt COPY . . -CMD ["python", "server.py"] +ENTRYPOINT ["python", "server.py"] diff --git a/README.md b/README.md index f79b7a1..7fc6b7f 100644 --- a/README.md +++ b/README.md @@ -33,15 +33,26 @@ The server runs at `http://127.0.0.1:5000`. In one terminal, connect to the WebSocket tunnel: -```powershell -websocat ws://127.0.0.1:5000/tunnel +```sh +websocat ws://127.0.0.1:5000/tunnel/asdf + +# with an api key +websocat \ + --header="X-API-Key:hello" \ + - ws://127.0.0.1:5000/tunnel/asdf ``` In another terminal, send a webhook payload: ### bash ```sh -curl -X POST http://127.0.0.1:5000/webhook \ +curl -X POST http://127.0.0.1:5000/webhook/asdf \ + -H "Content-Type: application/json" \ + -d '{"message":"hello from webhook"}' + +# with an api key +curl -X POST http://127.0.0.1:5000/webhook/asdf \ + -H "X-API-Key: hello" \ -H "Content-Type: application/json" \ -d '{"message":"hello from webhook"}' ``` @@ -52,3 +63,30 @@ Invoke-RestMethod -Method Post -Uri http://127.0.0.1:5000/webhook -ContentType " ``` The connected WebSocket client should receive the JSON payload. + +## on the real website +create an file called `config.yml` with content like +```yml +api_key: thesecretofalltime +``` +run the server with +```sh +docker-compose up --build -d +``` + +if you wanna see logs +```sh +docker logs smee2-smee2-1 --tail 300 -f +``` + +to test that its working, use the same api key like: + +```sh +curl -X POST https://sce.sjsu.edu/webhook/asdf \ + -H "Content-Type: application/json" \ + -d '{"message":"hello from webhook"}' + +websocat \ + --header="X-API-Key:hello" \ + - wss://sce.sjsu.edu/tunnel/asdf +``` diff --git a/args.py b/args.py index bf5e48f..0f2693a 100644 --- a/args.py +++ b/args.py @@ -8,7 +8,8 @@ def get_args(): "-v", action="count", default=0, - help="increase logging verbosity; can be used multiple times" + help="increase logging verbosity; can be used multiple times like -vvv" ) + parser.add_argument("--config", help="path to yaml file for api key, see readme") return parser.parse_args() diff --git a/docker-compose.yml b/docker-compose.yml index a625591..e35dd15 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -3,6 +3,12 @@ services: build: context: . dockerfile: ./Dockerfile - ports: - - "5000:5000" + volumes: + - ./config.yml:/app/config.yml + command: + - --config=/app/config.yml +networks: + default: + external: + name: sce diff --git a/requirements.txt b/requirements.txt index 6bd4692..fc1c223 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,4 +18,3 @@ typing_extensions==4.15.0 uvicorn==0.49.0 watchfiles==1.2.0 websockets==16.0 - diff --git a/server.py b/server.py index 4867a7e..d5b4e71 100644 --- a/server.py +++ b/server.py @@ -1,9 +1,11 @@ import collections import logging +import collections from fastapi import FastAPI, WebSocket, Request, HTTPException, Response import prometheus_client import uvicorn +import yaml from args import get_args from metrics import MetricsHandler @@ -19,6 +21,8 @@ ) logging.getLogger("uvicorn.access").setLevel(logging.WARNING) logging.getLogger("uvicorn.error").setLevel(logging.WARNING) +logging.getLogger("asyncio").setLevel(logging.WARNING) +logging.getLogger("watchfiles").setLevel(logging.WARNING) logger = logging.getLogger(__name__) @@ -27,35 +31,45 @@ clients = collections.defaultdict(list) subscribers = {} +# see if __name__ == '__server__' section for setting this value +API_KEY = None @app.post("/webhook/{subscription_id}") async def webhook(subscription_id: str, request: Request): - if subscription_id is not None: - header_val = request.headers.get("X-API-Key") - if (header_val != "hello"): - raise HTTPException(status_code=403, detail="API key is not valid ") - - data = await request.json() - logger.debug("Data pushed to webhook %s, received: %s", subscription_id, data) - - subscribers[subscription_id] = data - - for client in clients.get(subscription_id, []): - await client.send_json(data) - - logger.error("Data sent to websocket client") - return {"message":"received"} - - else: - logger.error("Invalid subscription '%s', connection not accepted", subscription_id) - return - - + if subscription_id is None: + logger.error( + "Invalid subscription '%s', connection not accepted", subscription_id + ) + MetricsHandler.failed_connections.labels( + subscription_id=subscription_id, + reason="no_subscription_id", + ).inc() + raise HTTPException( + status_code=404, + detail="subscripiton id was empty. expecting soemthing like /webhook/1234", + ) + header_val = request.headers.get("X-API-Key") + + data = await request.json() + logger.debug("Data pushed to webhook %s, received: %s", subscription_id, data) + + subscribers[subscription_id] = data + + for client in clients.get(subscription_id, []): + await client.send_json(data) + + number_of_clients = len(clients.get(subscription_id, [])) + message = f"forwarded to {number_of_clients} client(s) for subscription {subscription_id}" + logger.debug(message) + return {"message": message} + + @app.websocket("/tunnel/{subscription_id}") async def websocket_endpoint(subscription_id: str, websocket: WebSocket): - api_key = websocket.headers.get("X-API-Key") - - if (api_key != "hello"): + api_key_from_header = websocket.headers.get("X-API-Key") + + if API_KEY is not None and api_key_from_header != API_KEY: + logger.error('/tunnel recieved invalid X-API-Key of "%s"', api_key_from_header) MetricsHandler.failed_connections.labels( subscription_id=subscription_id, reason="bad_api_key", @@ -64,29 +78,33 @@ async def websocket_endpoint(subscription_id: str, websocket: WebSocket): return await websocket.accept() - + MetricsHandler.connected_clients.labels(subscription_id).inc() - + clients[subscription_id].append(websocket) - logger.debug(f"Websocket connection successfully established at id: {subscription_id}") - + logger.debug( + f"Websocket connection successfully established at id: {subscription_id}" + ) + try: while True: data = await websocket.receive_text() await websocket.send_text("Message received") - except Exception as e: + except Exception: + logger.exception("ok") MetricsHandler.failed_connections.labels( subscription_id=subscription_id, reason="websocket_receive_failed", ).inc() MetricsHandler.connected_clients.labels(subscription_id).dec() + finally: clients[subscription_id].remove(websocket) if not clients[subscription_id]: clients.pop(subscription_id, None) logger.debug(f"Websocket connection disconnected at id: {subscription_id}") - + @app.get("/metrics") def get_metrics(): @@ -95,6 +113,7 @@ def get_metrics(): media_type="text/plain", ) + # we have a separate __name__ check here due to how FastAPI starts # a server. the file is first ran (where __name__ == "__main__") # and then calls `uvicorn.run`. the call to run() reruns the file, @@ -105,6 +124,19 @@ def get_metrics(): # server uses if __name__ == "server": MetricsHandler.init() + try: + with open(args.config, "r") as stream: + data = yaml.safe_load(stream) + API_KEY = data.get("api_key", None) + logger.info(f'loaded api key from {args.config}') + except Exception: + logging.warning("unable to open yaml file, smee2 is not checking for api keys") if __name__ == "__main__": - uvicorn.run("server:app", host="0.0.0.0", port=5000) + uvicorn.run( + "server:app", + host="0.0.0.0", + port=5000, + reload=True, + timeout_graceful_shutdown=1, + )