import asyncio
import uuid
import websockets
from websockets.server import serve
from mythic_container.logging import logger
from mythic_container.grpc.pushC2GRPC_pb2_grpc import PushC2Stub
from mythic_container.grpc import pushC2GRPC_pb2 as grpcFuncs
import grpc.aio
import json
UUIDToWebsocketConn = {}
grpcStream = None
async def handleStreamConnection(client):
global UUIDToWebsocketConn
global grpcStream
try:
while True:
grpcStream = client.StartPushC2StreamingOneToMany()
await grpcStream.write(grpcFuncs.PushC2MessageFromAgent(
C2ProfileName="websocket"
))
logger.info(f"Connected to gRPC for pushC2 StreamingOneToMany")
async for request in grpcStream:
# this is streaming responses from Mythic to go to agents
try:
if request.TrackingID in UUIDToWebsocketConn:
logger.info(f"sending message back to websocket for id: {request.TrackingID}")
await UUIDToWebsocketConn[request.TrackingID].send(json.dumps({"data": request.Message.decode()}))
else:
logger.error(f"tracking ID not tracked: {request.TrackingID} ")
except Exception as d:
logger.exception(f"Failed to process handleStreamConnection message:\n{d}")
logger.error(f"disconnected from gRPC for handleStreamConnection")
except Exception as e:
logger.exception(f"[-] exception in handleStreamConnection: {e}")
async def handleGrpcStreamingServices():
maxInt = 2 ** 31 - 1
while True:
try:
logger.info(f"Attempting connection to gRPC for pushC2OneToMany...")
channel = grpc.aio.insecure_channel(
f'127.0.0.1:17444',
options=[
('grpc.max_send_message_length', maxInt),
('grpc.max_receive_message_length', maxInt),
])
await channel.channel_ready()
client = PushC2Stub(channel=channel)
streamConnections = handleStreamConnection(client)
logger.info(f"[+] Successfully connected to gRPC for pushC2OneToMany")
await asyncio.gather(streamConnections)
except Exception as e:
logger.exception(f"Translation gRPC services closed for pushC2OneToMany: {e}")
async def handle_connection(websocketConn: websockets.WebSocketServerProtocol):
global UUIDToWebsocketConn
global grpcStream
connUUID = str(uuid.uuid4())
logger.info(f"New tracking ID created: {connUUID}")
UUIDToWebsocketConn[connUUID] = websocketConn
try:
async for message in websocketConn:
# get message from agent and send it to grpc stream
logger.info(f"new websocket msg for id: {connUUID}")
while True:
if grpcStream is None:
await asyncio.sleep(1)
continue
break
try:
jsonMsg = json.loads(message)
await grpcStream.write(grpcFuncs.PushC2MessageFromAgent(
C2ProfileName="websocket",
RemoteIP=str(websocketConn.remote_address),
Base64Message=jsonMsg["data"].encode("utf-8"),
TrackingID=connUUID
))
except Exception as e:
logger.info(f"Hit exception trying to send websocket message to grpc: {e}")
await asyncio.sleep(1)
except Exception as c:
if grpcStream is not None:
logger.info(f"websocket connection dead, removing it: {connUUID}")
try:
del UUIDToWebsocketConn[connUUID]
await grpcStream.write(grpcFuncs.PushC2MessageFromAgent(
C2ProfileName="websocket",
RemoteIP=str(websocketConn.remote_address),
TrackingID=connUUID,
AgentDisconnected=True
))
except Exception as e:
logger.error(f"Failed to send message to Mythic that connection dropped: {e}")
async def main():
logger.info("starting grpc connection server")
asyncio.create_task(handleGrpcStreamingServices())
logger.info("starting websocket server")
async with serve(handle_connection, "127.0.0.1", 8081):
await asyncio.Future()
asyncio.run(main())