add put/get noise for agent functionality, refactor streaming

Author: st0ffregen

checker/src/checker.py

@@ -99,6 +99,16 @@ async def put(self, url, json=None, headers=None):
     async def get(self, url, headers=None):
         return await self.client.post(url, headers=self.apply_headers(headers))
 
+    async def stream(self, url, json=None, headers=None):
+        async with self.client.stream("POST", url, json=json, headers=self.apply_headers(headers)) as resp:
+            resp.raise_for_status()
+            try:
+                async for _ in resp.aiter_lines():
+                    break
+            except:
+                pass
+            return resp
+
     async def register_user(self, username: str, password: str) -> Tuple[int, int]:
         self.logger.debug(f"Registering user {username}")
         try:
@@ -299,6 +309,20 @@ async def mcp_read_messages(self, chat_session_id: int, token: str = None) -> li
             self.logger.error(f"Error during reading messages through MCP of session id {chat_session_id}: {e}")
             raise MumbleException("Error during reading messages through MCP.")
 
+    async def send_stream(self, message: str, chat_session_id: int):
+        self.logger.debug(f"Sending stream message to chat session {chat_session_id}")
+        try:
+            resp = await self.stream(f"{API_BASE}/messages/stream", json={"input": message, "ChatSessionId": chat_session_id, "SkipLlmResponse": True})
+        except (ConnectTimeout, NetworkError, PoolTimeout, RequestError) as e:
+            self.logger.error(f"Connection error during streaming message to chat session {chat_session_id}: {e}")
+            raise MumbleException(f"Connection error during streaming message.")
+        except HTTPStatusError as e:
+            self.logger.error(f"HTTP error during streaming message to chat session {chat_session_id}: {e}")
+            raise MumbleException(f"HTTP error during streaming message.")
+        except Exception as e:
+            self.logger.error(f"Error during streaming message to chat session {chat_session_id}: {e}")
+            raise MumbleException(f"Error during streaming message.")
+
 
 # Dependency injection for Connection
 @checker.register_dependency
@@ -730,21 +754,7 @@ async def putflag_agent(
     flag = task.flag
     flag = await format_flag(flag)
 
-    async def send_stream(conn, flag, chat_session_id):
-        async with conn.client.stream(
-                "POST",
-                f"{API_BASE}/messages/stream",
-                json={"input": flag, "ChatSessionId": chat_session_id, "SkipLlmResponse": True},
-                headers=conn.headers
-        ) as resp:
-            resp.raise_for_status()
-            try:
-                async for _ in resp.aiter_lines():
-                    break
-            except:
-                pass
-
-    await send_stream(conn, flag, chat_session_id)
+    await conn.send_stream(flag, chat_session_id)
 
     access_token = await conn.create_access_token(random_string(4))
 
@@ -877,6 +887,91 @@ async def exploit_agent(
     return unformatted_flag
 
 
+@checker.putnoise(1)
+async def putnoise_agent(
+    conn: Connection,
+    db: ChainDB,
+):
+    username = random_string()
+    password = random_string()
+    user_id, chat_session_id = await conn.register_user(username, password)
+    await conn.login_user(username, password)
+    
+    # Generate random noise sentence
+    noise_sentence = ''.join(random.choices(get_ascii_chars(), k=SUPPOSED_FLAG_LENGTH))
+
+    await conn.send_stream(noise_sentence, chat_session_id)
+
+    access_token = await conn.create_access_token(random_string(4))
+
+    await db.set("userdata", {"user_id": user_id, "session_id": chat_session_id, "access_token": access_token, "noise_sentence": noise_sentence})
+
+
+@checker.getnoise(1)
+async def getnoise_agent(
+    task: GetflagCheckerTaskMessage, conn: Connection, db: ChainDB, logger: LoggerAdapter
+) -> None:
+
+    try:
+        data = await db.get("userdata")
+        user_id = data["user_id"]
+        access_token = data["access_token"]
+        session_id = data["session_id"]
+        noise_sentence = data["noise_sentence"]
+    except KeyError:
+        logger.error("Missing database entry from putnoise(1)")
+        raise MumbleException("Missing database entry from putnoise(1)")
+
+    client_challenge_bytes = os.urandom(8)
+    client_challenge_json = base64.b64encode(client_challenge_bytes).decode('utf-8')
+    server_challenge, memory_key, iv = await conn.client_challenge(client_challenge_json, user_id)
+
+    salt = client_challenge_bytes + server_challenge
+    session_key = HKDF(access_token, 16, salt, SHA256, 1)
+
+    client_credentials_bytes = aes_cfb8_encrypt(client_challenge_bytes, session_key, iv)
+    client_credentials_json = base64.b64encode(client_credentials_bytes).decode('utf-8')
+    agent_name = random_string(4)
+    try:
+        server_creds, jwt = await conn.client_credentials(client_credentials_json, user_id, agent_name, memory_key)
+    except (ConnectTimeout, NetworkError, PoolTimeout, RequestError) as e:
+        logger.error(
+            f"Connection error during posting client credentials for user {user_id} for agent {agent_name} with memory key {memory_key} and client credentials {client_credentials_json}: {e}")
+        raise MumbleException(
+            f"Connection error during posting client credentials.")
+    except HTTPStatusError as e:
+        logger.error(
+            f"HTTP error during posting client credentials for user {user_id} for agent {agent_name} with memory key {memory_key} and client credentials {client_credentials_json}: {e}")
+        raise MumbleException("HTTP error during posting client credentials.")
+    except Exception as e:
+        logger.error(
+            f"Error during posting client credentials for user {user_id} for agent {agent_name} with memory key {memory_key} and client credentials {client_credentials_json}: {e}")
+        raise MumbleException("Error during posting client credentials.")
+
+    server_creds_computed = aes_cfb8_encrypt(server_challenge, session_key, iv)
+
+    if server_creds_computed != server_creds:
+        logger.error(f"Server credentials do not match. Computed ones: {server_creds_computed}. Actual {server_creds}")
+        raise MumbleException("Server credentials do not match")
+
+    try:
+        mcp_client = MCPClient(f"http://{task.address}:{SERVICE_PORT}/mcp", logger, jwt)
+        messages = await mcp_client.run(session_id)
+        assert len(messages) >= 2, "Expected at least 2 messages from MCP"
+        retrieved_sentence = messages[1]["content"]
+    except Exception as e:
+        logger.error(f"Error during reading messages through MCP of session id {session_id}: {e}")
+        raise MumbleException(f"Error during reading messages through MCP.")
+
+    if retrieved_sentence is None or retrieved_sentence == "":
+        logger.error("No able to retrieve noise sentence")
+        raise MumbleException("No able to retrieve noise sentence")
+
+    if noise_sentence != retrieved_sentence:
+        logger.error(f"Noise sentence mismatch: expected {noise_sentence}, got {retrieved_sentence}")
+        raise MumbleException("Noise sentence does not match expected value")
+
+
 if __name__ == "__main__":
 
     debug = False