user handle, round-robin balancer

.gitignore

+# Python bytecode:
+*.py[co]
+.tokens
+
+# Packaging files:
+*.egg*
+
+# Editor temp files:
+*.swp
+*.swo
+*~
+.vscode
+.idea
+
+# Sphinx docs:
+build
+_build
+
+# SQLite3 database files:
+*.db
+
+# Logs:
+*.log
+.ropeproject
+celerybeat-schedule
+.coverage
+*,cover
+coverage.xml
+.noseids

logging.conf

+[loggers]
+keys=root,uicheckapp
+
+[handlers]
+keys=consoleHandler,detailedConsoleHandler
+
+[formatters]
+keys=normalFormatter,detailedFormatter
+
+[logger_root]
+level=INFO
+handlers=consoleHandler
+
+[logger_uicheckapp]
+level=DEBUG
+handlers=detailedConsoleHandler
+qualname=uicheckapp
+propagate=0
+
+[handler_consoleHandler]
+class=StreamHandler
+level=DEBUG
+formatter=normalFormatter
+args=(sys.stdout,)
+
+[handler_detailedConsoleHandler]
+class=StreamHandler
+level=DEBUG
+formatter=detailedFormatter
+args=(sys.stdout,)
+
+[formatter_normalFormatter]
+format=%(asctime)s loglevel=%(levelname)-6s logger=%(name)s %(funcName)s() L%(lineno)-4d %(message)s
+
+[formatter_detailedFormatter]
+format=%(asctime)s loglevel=%(levelname)-6s logger=%(name)s %(funcName)s() L%(lineno)-4d %(message)s   call_trace=%(pathname)s L%(lineno)-4d
\ No newline at end of file

main.py

+from fastapi import FastAPI, Header
+from fastapi.responses import ORJSONResponse
+import json
+from typing import Union, List
+import requests
+from sredis.sredis import * 
+import logging
+from models import User, DataCenter, Token
+from uuid import uuid4
+
+logging.config.fileConfig('logging.conf', disable_existing_loggers=False)
+
+# get root logger
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+add_datacenter("https://kappa1.fured.cloud.bme.hu")
+add_datacenter("https://kappa2.fured.cloud.bme.hu")
+add_datacenter("https://kappa3.fured.cloud.bme.hu")
+
+def _proxy_datacenters(serverpath: str, username, method="GET", balancer_fun = rr_get):
+    server = balancer_fun()
+    token = get_datacenter_token(username, server)
+    url=f"{server}/{serverpath}"
+    logger.debug("Req: " + url)
+    t_resp = requests.request(
+        method=method,
+        url=url,
+        allow_redirects=False, verify=False, 
+        headers={
+            'Authorization': token
+        }
+    )
+    
+    response = ORJSONResponse(
+        json.loads(t_resp.content), 
+        status_code=t_resp.status_code
+    )
+    return response
+
+@app.get("/lb/{server_path:path}")
+def proxy(
+        Authorization: Union[str, None] = Header(default=None), 
+        server_path: str = "/"
+    ):
+    username = get_user(str(Authorization))
+    if not username:
+        return Response(status_code=401)
+    
+    return _proxy_datacenters(server_path, username)
+
+@app.post("/lb/{server_path:path}")
+def proxy(
+        Authorization: Union[str, None] = Header(default=None),
+        server_path: str = "/"
+    ):
+    username = get_user(str(Authorization))
+    if not username:
+        return Response(status_code=401)
+    return _proxy_datacenters(server_path, username, method="POST")
+
+@app.post("/add_datacenter/")
+def create_datacenter(
+        Authorization: Union[str, None] = Header(default=None),
+        dc: DataCenter = None
+    ):
+    username = get_user(str(Authorization))
+    if not username:
+        return Response(status_code=401)
+    add_datacenter(dc.name)
+    return Response(status_code=201)
+
+
+@app.post("/create_user/")
+def create_user(
+        user: User = None
+    ):
+    user.token = str(uuid4())
+    new_user(user.username, user.token)
+    return user
+
+@app.post("/set_tokens/")
+def set_tokens(
+        Authorization: Union[str, None] = Header(default=None),
+        tokens: List[Token] = None
+    ):
+    username = get_user(str(Authorization))
+    if not username:
+        return Response(status_code=401)
+    for token in tokens:
+        set_token(username, str(token.datacenter), str(token.token))
+    return tokens

models.py

+from pydantic import BaseModel
+from typing import Union
+
+class User(BaseModel):
+    username: str
+    token: Union[str, None] = None
+
+class DataCenter(BaseModel):
+    name: str
+
+class Token(BaseModel):
+    datacenter: str
+    token: str
\ No newline at end of file

pyproject.toml

+[tool.poetry]
+name = "balancer"
+version = "0.1.0"
+description = ""
+authors = ["Karsa Zoltán István <karsazoltan@gmail.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.7"
+fastapi = "^0.93.0"
+poethepoet = "^0.18.1"
+uvicorn = "^0.20.0"
+requests = "^2.28.2"
+redis = {extras = ["hiredis"], version = "^4.5.1"}
+orjson = "^3.8.7"
+
+[tool.poe.tasks.start]
+shell = "poetry run uvicorn main:app --reload --port 6973 --host 0.0.0.0"
+help   = "Start the microservice on port 6973"
+
+[tool.poe.tasks.lint]
+shell = "poetry run black . && poetry run ruff --fix . && poetry run mypy backend"
+help = "Lint the most important parts of the microservice with black"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

sredis/sredis.py

+import redis
+
+r = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
+r.set("datacenters_cnt", 1)
+all_keys = list(r.hgetall('datacenters_hash').keys())
+if all_keys:
+    r.hdel('datacenters_hash', *all_keys)
+
+def add_datacenter(datacenter: str):
+    cnt = int(r.get("datacenters_cnt"))
+    r.hset("datacenters_hash", cnt, datacenter)
+    r.incr("datacenters_cnt")
+    r.set("roundrobin_cnt", 1)
+
+def rr_get():
+    cnt = int(r.get("datacenters_cnt"))
+    rr  = int(r.get("roundrobin_cnt"))
+    if rr + 1 >= cnt:
+        r.set("roundrobin_cnt", 1)
+    else:
+        r.incr("roundrobin_cnt")
+    return str(r.hget("datacenters_hash", rr))
+
+def new_user(username: str, token: str):
+    if not r.hexists("users", token):
+        r.hset("users", token, username)
+
+def get_user(token: str):
+    user = str(r.hget("users", token))
+    return user
+
+def set_token(username: str, datacenter: str, token: str):
+    r.hset(f"tokens:{username}", datacenter, token)
+
+def get_datacenter_token(username: str, datacenter: str):
+    return str(r.hget(f"tokens:{username}", datacenter))
\ No newline at end of file