Reafctor app

app/util/websocket_auth_backend.py

@@ -0,0 +1,38 @@
+import base64
+import binascii
+from uuid import UUID
+from starlette.authentication import (
+    AuthCredentials, AuthenticationBackend, AuthenticationError, SimpleUser
+)
+
+from app.database import SessionLocal
+from app.models.chargepoint import ChargePoint
+
+class BasicAuthBackend(AuthenticationBackend):
+    async def authenticate(self, conn):
+        if "Authorization" not in conn.headers:
+            raise AuthenticationError('No Authorization header provided')
+
+        auth = conn.headers["Authorization"]
+        try:
+            scheme, credentials = auth.split()
+            if scheme.lower() != 'basic':
+                raise AuthenticationError('Invalid authorization scheme')
+            decoded = base64.b64decode(credentials).decode("ascii")
+        except (ValueError, UnicodeDecodeError, binascii.Error) as exc:
+            raise AuthenticationError('Invalid basic auth credentials')
+
+        username, _, password = decoded.partition(":")
+        try:
+            id = UUID(username)
+        except (ValueError) as exc:
+            raise AuthenticationError('Invalid basic auth credentials')
+        
+        with SessionLocal() as db:
+            chargepoint = db.query(ChargePoint).filter(ChargePoint.id == id).first()
+            if chargepoint is None:
+                raise AuthenticationError('Invalid basic auth credentials')
+            if chargepoint.password != password:
+                raise AuthenticationError('Invalid basic auth credentials')
+            
+        return AuthCredentials(["authenticated"]), chargepoint

app/util/websocket_wrapper.py

@@ -0,0 +1,19 @@
+from fastapi import WebSocket, WebSocketDisconnect
+from websockets import ConnectionClosed
+
+# Wrapper to transform a FastAPI websocket to a standard websocket
+class WebSocketWrapper():
+    def __init__(self, websocket: WebSocket):
+        self._websocket = websocket
+
+    async def recv(self) -> str:
+        try:
+            return await self._websocket.receive_text()
+        except WebSocketDisconnect as e:
+            raise ConnectionClosed(e.code, 'WebSocketWrapper')
+
+    async def send(self, msg: str) -> None:
+        await self._websocket.send_text(msg)
+
+    async def close(self, code: int, reason: str) -> None:
+        await self._websocket.close(code)