add: socket io authetication enforcing

backend/app.py

@@ -30,9 +30,8 @@
 from models.response import MessageResponse, MessageResponseInvalidError, ResponseStatus
 from typing import Any
 from models.config import Configuration, SetupStatus, StatusAPI
-from utils import json_like
-from utils import crypto
-
+from utils import json_like, crypto
+from utils.auth import login_validation, AuthStatus
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
 os.makedirs(EXPLOIT_SOURCES_DIR, exist_ok=True)
 
@@ -67,23 +66,17 @@ async def create_access_token(data: dict) -> str:
     return encoded_jwt
 
 async def is_loggined(token: str = Depends(oauth2_scheme)) -> None|bool:
-
-    #App status checks
-    config = await Configuration.get_from_db()
-
-    if not config.login_enabled:
-        return True
-
-    #If the app is running and requires login
-    if not token:
-        return None
-
-    try:
-        payload = jwt.decode(token, await APP_SECRET(), algorithms=[JWT_ALGORITHM])
-        authenticated: bool = payload.get("authenticated", False)
-        return authenticated
-    except Exception:
-        return False
+    result = await login_validation(token)
+    match result:
+        case AuthStatus.ok:
+            return True
+        case AuthStatus.nologin:
+            return None
+        case AuthStatus.invalid:
+            return False
+        case AuthStatus.wrong:
+            return False
+    return None
 
 async def check_login(status: bool|None = Depends(is_loggined)) -> None:
     if status is None:

backend/db.py

@@ -48,8 +48,9 @@ class redis_channels:
     submitter = "submitter"
     stats = "stats"
     config = "config"
+    password_change = "password_change"
 
-REDIS_CHANNEL_LIST = [
+REDIS_CHANNEL_PUBLISH_LIST = [
     "client",
     "attack_group",
     "exploit",
@@ -290,6 +291,7 @@ async def regen_app_secret():
             )
         )
         await redis_conn.delete("APP_SECRET")
+        await redis_conn.publish(redis_channels.password_change, "update")
 
 APP_SECRET = get_dbenv_func("APP_SECRET", new_app_secret, value_cached=True)
 SERVER_ID = get_dbenv_func("SERVER_ID", uuid4, value_cached=True)
@@ -340,7 +342,6 @@ async def _dbtransaction():
         yield session
         await session.commit()
 
-
 DBSession = Annotated[AsyncSession, Depends(_dbtransaction)]
 
 async def connect_db():

backend/models/enums.py

@@ -9,6 +9,12 @@ class ExploitStatus(Enum):
     active = 'active'
     disabled = 'disabled'
 
+class AuthStatus(Enum):
+    ok = "ok"
+    nologin = "nologin"
+    wrong = "wrong"
+    invalid = "invalid"
+
 class FlagStatus(Enum):
     ok = 'ok'
     wait = 'wait'

backend/skio.py

@@ -5,11 +5,12 @@
 import logging
 from env import DEBUG
 from multiprocessing import Process
-from db import close_db, redis_conn, REDIS_CHANNEL_LIST, connect_db, dbtransaction, redis_channels
+from db import close_db, redis_conn, REDIS_CHANNEL_PUBLISH_LIST, connect_db, dbtransaction, redis_channels
 from utils.query import get_exploits_with_latest_attack, detailed_exploit_status
 from models.config import Configuration
 from models.enums import ExploitStatus
-
+from utils.auth import login_validation, AuthStatus
+from socketio.exceptions import ConnectionRefusedError
 class StopLoop(Exception):
     pass
 
@@ -26,23 +27,51 @@ class StopLoop(Exception):
 class g:
     task_list = []
 
+async def check_login(token: str) -> None:
+    status = await login_validation(token)
+    if status == AuthStatus.ok:
+        return None
+    if status == AuthStatus.nologin:
+        raise ConnectionRefusedError("Authentication required")
+    raise ConnectionRefusedError("Unauthorized")
+
 @sio_server.on("connect")
-async def sio_connect(sid, environ): pass
+async def sio_connect(sid, environ, auth):
+    await check_login(auth.get("token"))
+    await redis_conn.lpush("sid_list", sid)
 
 @sio_server.on("disconnect")
-async def sio_disconnect(sid): pass
+async def sio_disconnect(sid):
+    await redis_conn.lrem("sid_list", 0, sid)
+
+async def disconnect_all():
+    while True:
+        sids = await redis_conn.lpop("sid_list", count=100)
+        if sids is None or len(sids) == 0:
+            break
+        for sid in sids:
+            await sio_server.disconnect(sid)
 
 async def generate_listener_tasks():
-    for chann in REDIS_CHANNEL_LIST:
+    for chann in REDIS_CHANNEL_PUBLISH_LIST:
         async def listener(chann=chann):
             await sio_server.emit(chann, "init")
             async with redis_conn.pubsub() as pubsub:
                 await pubsub.subscribe(chann)
                 while True:
                     message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=None)
-                    await sio_server.emit(chann, message)
+                    if message:
+                        await sio_server.emit(chann, message)
         g.task_list.append(asyncio.create_task(listener()))
 
+async def password_change_listener():
+    async with redis_conn.pubsub() as pubsub:
+        await pubsub.subscribe(redis_channels.password_change)
+        while True:
+            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=None)
+            if message:
+                await disconnect_all()
+
 async def check_exploits_disabled():
     disabled_exploits = set([])
     async with dbtransaction() as db:
@@ -67,9 +96,10 @@ async def tasks_init():
     try:
         await connect_db()
         await generate_listener_tasks()
-        check_disabled_exploits = asyncio.create_task(check_exploits_disabled())
+        pwd_change = asyncio.create_task(password_change_listener())
+        check_disab = asyncio.create_task(check_exploits_disabled())
         logging.info("SocketIO manager started")
-        await asyncio.gather(*g.task_list, check_disabled_exploits)
+        await asyncio.gather(*g.task_list, check_disab, pwd_change)
     except KeyboardInterrupt:
         pass
     finally:

backend/utils/__init__.py

@@ -11,6 +11,7 @@
 import logging
 from env import EXPLOIT_SOURCES_DIR
 
+
 #logging.getLogger().setLevel(logging.DEBUG)
 logging.basicConfig(format="[EXPLOIT-FARM][%(asctime)s] >> [%(levelname)s][%(name)s]:\t%(message)s", datefmt="%d/%m/%Y %H:%M:%S")
 crypto = CryptContext(schemes=["bcrypt"], deprecated="auto")
@@ -173,5 +174,4 @@ async def pubsub_flush(pubsub):
     while flushed is not None:
         flushed = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0)
 
-
 

backend/utils/auth.py

@@ -0,0 +1,27 @@
+from models.config import Configuration
+from db import APP_SECRET
+from env import JWT_ALGORITHM
+import jwt
+from models.enums import AuthStatus
+
+async def login_validation(token: str|None) -> AuthStatus:
+
+    #App status checks
+    config = await Configuration.get_from_db()
+
+    if not config.login_enabled:
+        return AuthStatus.ok
+
+    #If the app is running and requires login
+    if not token:
+        return AuthStatus.nologin
+
+    try:
+        payload = jwt.decode(token, await APP_SECRET(), algorithms=[JWT_ALGORITHM])
+        authenticated: bool = payload.get("authenticated", False)
+        if authenticated:
+            return AuthStatus.ok
+        else:
+            return AuthStatus.wrong
+    except Exception:
+        return AuthStatus.invalid

client/exploitfarm/utils/config.py

@@ -41,7 +41,12 @@ def skio_connect(self) -> socketio.Client:
         from exploitfarm import DEV_MODE
         from exploitfarm.utils.reqs import get_url, ReqsError
         try:
-            self.skio.connect(get_url("//", self), socketio_path="/sock/socket.io", transports=["websocket"])
+            self.skio.connect(
+                get_url("//", self),
+                socketio_path="/sock/socket.io",
+                transports=["websocket"],
+                auth={"token": self.server.auth_key}
+            )
         except Exception as e:
             if DEV_MODE:
                 traceback.print_exc()

client/exploitfarm/xploit.py

@@ -261,9 +261,6 @@ class InvalidSploitError(Exception):
     fdlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
     resource.setrlimit(resource.RLIMIT_NOFILE, (fdlimit[1]//2,fdlimit[1]))
 
-
-
-
 class APIException(Exception):
     pass
 

frontend/src/App.tsx

@@ -8,7 +8,7 @@ import { notifications, Notifications } from '@mantine/notifications';
 import { LoadingOverlay, MantineProvider, Title } from '@mantine/core';
 import { LoginProvider } from '@/components/LoginProvider';
 import { Routes, Route, BrowserRouter } from "react-router-dom";
-import { useGlobalStore } from './utils/stores';
+import { useGlobalStore, useTokenStore } from './utils/stores';
 import { statusQuery } from './utils/queries';
 import { HomePage } from './components/screens/HomePage';
 import { MainLayout } from './components/MainLayout';
@@ -20,7 +20,8 @@ import { useDebouncedCallback } from '@mantine/hooks';
 export default function App() {
 
     const queryClient = useQueryClient()
-    const { setErrorMessage } = useGlobalStore()
+    const { setErrorMessage, loading:loadingStatus } = useGlobalStore()
+    const { loginToken } = useTokenStore()
 
     const debouncedCalls = DEBOUNCED_SOCKET_IO_CHANNELS.map((channel) => (
         useDebouncedCallback(() => {
@@ -49,6 +50,7 @@ export default function App() {
                 color: "red"
             })
         });
+
         let first_time = true
         socket_io.on("connect", () => {
             if (socket_io.connected) {
@@ -62,7 +64,6 @@ export default function App() {
                         color: "blue",
                         icon: "🚀",
                     })
-
                 }
             }
             first_time = false
@@ -75,7 +76,12 @@ export default function App() {
         }
     }, [])
 
-    const loadingStatus = useGlobalStore((store) => store.loading)
+
+    useEffect(() => {
+        socket_io.auth = { token: loginToken }
+        socket_io.disconnect()
+        socket_io.connect()
+    }, [loginToken])
 
     const status = statusQuery()
 

frontend/src/components/screens/SetupScreen.tsx

@@ -235,13 +235,13 @@ export const SetupScreen = ({ editMode, onSubmit }:{ editMode?:boolean, onSubmit
         <Divider my="md" />
         <Title order={2}><u>Authentication</u></Title>
         <Space h="md" />
-        {(editMode && !status.data?.config?.AUTHENTICATION_REQUIRED)?<>
+        {editMode?<>
             <Alert
                 color="yellow"
                 title="Warning"
                 icon={<TiWarning />}
             >
-                Enabling the authentication will trigger the stop of all the running exploits, you can run them again manually (the password will be required at that time)
+                Enabling the authentication or changing the password will trigger the stop of all the running exploits, you can run them again manually (the new password will be required at that time)
             </Alert>
             <Space h="md" />
         </>:null}

frontend/src/utils/net.ts

@@ -4,8 +4,20 @@ import io from 'socket.io-client';
 export const DEV_IP_BACKEND = "127.0.0.1:5050"
 
 export const socket_io = import.meta.env.DEV?
-    io("ws://"+DEV_IP_BACKEND, { path:"/sock/socket.io", transports: ['websocket'] }):
-    io({ path:"/sock/socket.io", transports: ['websocket'] })
+    io("ws://"+DEV_IP_BACKEND, {
+        path:"/sock/socket.io",
+        transports: ['websocket'],
+        auth: {
+            token: useTokenStore.getState().loginToken
+        }
+    }):
+    io({
+        path:"/sock/socket.io",
+        transports: ['websocket'],
+        auth: {
+            token: useTokenStore.getState().loginToken
+        }
+    })
 
 export const SOCKET_IO_CHANNELS = [
     "client",

frontend/src/utils/time.ts

@@ -33,7 +33,7 @@ export const durationToString = (duration: Duration): string => {
     if (duration.milliseconds() > 0 && duration.seconds() === 0) {
         result.push(`${duration.milliseconds()} ms`)
     }
-    if (result.length === 0) result.push(`0 ms`)
+    if (result.length === 0) result.push(`0 s`)
     return result.join(", ")
 }
 

tests/xploit_test/config.toml

@@ -3,4 +3,4 @@ name = "xploit_test"
 interpreter = "python3"
 run = "main.py"
 language = "python"
-service = "33f46bdc-464a-423d-9ea5-e1bdc14c98e1"
+service = "24b0547e-b327-4ca1-b251-cbd91fade18f"