def __init__(self):
self.users = Users()
self.loop = asyncio.get_running_loop()
self.contexts = {
'listeners': Listeners(self),
'sessions': Sessions(self),
'modules': Modules(self),
'stagers': Stagers(self),
'users': self.users
}
class TeamServer:
def __init__(self):
self.users = Users()
self.loop = asyncio.get_running_loop()
self.contexts = {
'listeners': Listeners(self),
'sessions': Sessions(self),
'modules': Modules(self),
'stagers': Stagers(self),
'users': self.users
}
async def process_client_message(self, user, path, data):
message = json.loads(data)
logging.debug(
f"Received message from {user.name}@{user.ip} path:{path} msg: {message}"
)
status = 'error'
try:
ctx = self.contexts[message['ctx'].lower()]
except KeyError:
traceback.print_exc()
result = f"Context '{message['ctx'].lower()}' does not exist"
logging.error(result)
else:
try:
cmd_handler = getattr(ctx, message['cmd'])
result = cmd_handler(**message['args'])
status = 'success'
except AttributeError:
traceback.print_exc()
result = f"Command '{message['cmd']}' does not exist in context '{message['ctx'].lower()}'"
except CmdError as e:
result = str(e)
except Exception as e:
traceback.print_exc()
result = f"Exception when executing command '{message['cmd']}': {e}"
logging.error(result)
await user.send({
'type': 'message',
'id': message['id'],
'ctx': message['ctx'],
'name': message['cmd'],
'status': status,
'result': result
})
async def update_server_stats(self):
stats = {
**{str(ctx): dict(ctx)
for ctx in self.contexts.values()}, 'ips': get_ips()
}
await self.users.broadcast_event(events.STATS_UPDATE, stats)
async def update_available_loadables(self):
loadables = {
str(ctx): [loadable.name for loadable in ctx.loaded]
for ctx in self.contexts.values() if hasattr(ctx, 'loaded')
}
await self.users.broadcast_event(events.LOADABLES_UPDATE, loadables)
async def connection_handler(self, websocket, path):
try:
user = await self.users.register(websocket)
await self.update_server_stats()
await self.update_available_loadables()
logging.info(f"New client connected {user.name}@{user.ip}")
except UsernameAlreadyPresentError as e:
logging.error(f"{websocket.remote_address[0]}: {e}")
return
while True:
try:
data = await asyncio.wait_for(websocket.recv(), timeout=20)
except asyncio.TimeoutError:
# No data in 20 seconds, check the connection.
logging.debug(
f"No data from {user.name}@{user.ip} after 20 seconds, sending ping"
)
try:
pong_waiter = await websocket.ping()
await asyncio.wait_for(pong_waiter, timeout=10)
except asyncio.TimeoutError:
# No response to ping in 10 seconds, disconnect.
logging.debug(
f"No pong from {user.name}@{user.ip} after 10 seconds, closing connection"
)
self.users.unregister(user.name)
await self.update_server_stats()
return
except websockets.exceptions.ConnectionClosed:
logging.debug(f"Connection closed by client")
self.users.unregister(user.name)
await self.update_server_stats()
return
else:
await self.process_client_message(user, path, data)