74 lines
No EOL
1.9 KiB
Python
74 lines
No EOL
1.9 KiB
Python
# session.py
|
|
|
|
import time
|
|
import asyncio
|
|
import websockets
|
|
|
|
import db
|
|
import mail
|
|
import agent
|
|
|
|
|
|
class Session(agent.Agent, mail.Mail):
|
|
"""Launch instance for every client connection.
|
|
|
|
Will recv messages from client and push them to mailbox.
|
|
Also, will pull messages from mailbox and send to client.
|
|
|
|
"""
|
|
def __init__(self, connection, crimata_id):
|
|
mail.Mail.__init__(self, crimata_id)
|
|
agent.Agent.__init__(self, connection)
|
|
|
|
self.on = True
|
|
|
|
print(f"Launched sesssion for {crimata_id}.")
|
|
|
|
# Recv intents run desired endpoint.
|
|
async def do_agent_intents(self):
|
|
while self.on:
|
|
|
|
try:
|
|
jpackage = await asyncio.wait_for(
|
|
self.recv(), timeout=0.1)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
except websockets.exceptions.ConnectionClosed:
|
|
await self.shutdown()
|
|
continue
|
|
|
|
# Check if target is verified.
|
|
verified = db.verify(jpackage.get("target"))
|
|
|
|
# Send error code back to agent.
|
|
if not verified:
|
|
await self.send(2)
|
|
continue
|
|
|
|
# Call corresponding endpoint.
|
|
await self.mailbox_send(jpackage)
|
|
|
|
# Send mail back to agent.
|
|
#! Pulling message from backend.
|
|
async def deliver_mail(self):
|
|
while self.on:
|
|
|
|
if not self.crimata_id:
|
|
await asyncio.sleep(0.1)
|
|
continue
|
|
|
|
# Try recv.
|
|
try:
|
|
jpackage = await asyncio.wait_for(
|
|
self.mailbox_recv(), timeout=0.1)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
# Send message to Agent.
|
|
print("Sending message")
|
|
await self.send(jpackage)
|
|
|
|
# Shutdown protocol.
|
|
async def shutdown(self):
|
|
print("Shutting down session.")
|
|
self.on = False |