bump v0.2.8-rc02

fixed bugs with how translation services handle timeouts
MythicMeta · May 12, 2023 · 795bb17 · 795bb17
1 parent 8af8233
commit 795bb17
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 57 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,11 @@
+
+## [v0.2.8-rc02] - 2023-05-12
+
+### Changed
+
+- Fixed some bugs with how translation services handle timeouts and reconnects
+
+
 ## [0.2.8-rc01] - 2023-05-10
 
 ### Changed

diff --git a/mythic_container/TranslationBase.py b/mythic_container/TranslationBase.py
@@ -101,71 +101,77 @@ async def handleTranslationServices(tr_name: str):
 
 async def handleGenerateKeys(tr_name: str, client):
     try:
-        stream = client.GenerateEncryptionKeys()
-        await stream.write(TrGenerateEncryptionKeysMessageResponse(
-            Success=True,
-            TranslationContainerName=tr_name
-        ))
-        logger.debug(f"Connected to gRPC for generating encryption keys for {tr_name}")
-        async for request in stream.__aiter__():
-            try:
-                result = await translationServices[tr_name].generate_keys(request)
-                result.TranslationContainerName = tr_name
-                await stream.write(result)
-            except Exception as d:
-                logger.exception(f"Failed to process message:\n{d}")
-                await stream.write(TrGenerateEncryptionKeysMessageResponse(
-                    Success=False,
-                    TranslationContainerName=tr_name,
-                    Error=f"Failed to process message:\n{d}"
-                ))
+        while True:
+            stream = client.GenerateEncryptionKeys()
+            await stream.write(TrGenerateEncryptionKeysMessageResponse(
+                Success=True,
+                TranslationContainerName=tr_name
+            ))
+            logger.info(f"Connected to gRPC for generating encryption keys for {tr_name}")
+            async for request in stream:
+                try:
+                    result = await translationServices[tr_name].generate_keys(request)
+                    result.TranslationContainerName = tr_name
+                    await stream.write(result)
+                except Exception as d:
+                    logger.exception(f"Failed to process message:\n{d}")
+                    await stream.write(TrGenerateEncryptionKeysMessageResponse(
+                        Success=False,
+                        TranslationContainerName=tr_name,
+                        Error=f"Failed to process message:\n{d}"
+                    ))
+            logger.error(f"disconnected from gRPC for generating encryption keys for {tr_name}")
     except Exception as e:
         logger.exception(f"[-] exception in handleGenerateKeys for {tr_name}")
 
 
 async def handleCustomToMythic(tr_name: str, client):
     try:
-        stream = client.TranslateFromCustomToMythicFormat()
-        await stream.write(TrCustomMessageToMythicC2FormatMessageResponse(
-            Success=True,
-            TranslationContainerName=tr_name
-        ))
-        logger.debug(f"Connected to gRPC for handling CustomC2 to MythicC2 Translations for {tr_name}")
-        async for request in stream.__aiter__():
-            try:
-                result = await translationServices[tr_name].translate_from_c2_format(request)
-                result.TranslationContainerName = tr_name
-                await stream.write(result)
-            except Exception as d:
-                logger.exception(f"Failed to process message:\n{d}")
-                await stream.write(TrCustomMessageToMythicC2FormatMessageResponse(
-                    Success=False,
-                    TranslationContainerName=tr_name,
-                    Error=f"Failed to process message:\n{d}"
-                ))
+        while True:
+            stream = client.TranslateFromCustomToMythicFormat()
+            await stream.write(TrCustomMessageToMythicC2FormatMessageResponse(
+                Success=True,
+                TranslationContainerName=tr_name
+            ))
+            logger.info(f"Connected to gRPC for handling CustomC2 to MythicC2 Translations for {tr_name}")
+            async for request in stream:
+                try:
+                    result = await translationServices[tr_name].translate_from_c2_format(request)
+                    result.TranslationContainerName = tr_name
+                    await stream.write(result)
+                except Exception as d:
+                    logger.exception(f"Failed to process message:\n{d}")
+                    await stream.write(TrCustomMessageToMythicC2FormatMessageResponse(
+                        Success=False,
+                        TranslationContainerName=tr_name,
+                        Error=f"Failed to process message:\n{d}"
+                    ))
+            logger.error(f"disconnected from gRPC for doing custom->mythic c2 for {tr_name}")
     except Exception as e:
         logger.exception(f"[-] exception in handleCustomToMythic for {tr_name}")
 
 
 async def handleMythicToCustom(tr_name: str, client):
     try:
-        stream = client.TranslateFromMythicToCustomFormat()
-        await stream.write(TrMythicC2ToCustomMessageFormatMessageResponse(
-            Success=True,
-            TranslationContainerName=tr_name
-        ))
-        logger.debug(f"Connected to gRPC for handling MythicC2 to CustomC2 Translations for {tr_name}")
-        async for request in stream.__aiter__():
-            try:
-                result = await translationServices[tr_name].translate_to_c2_format(request)
-                result.TranslationContainerName = tr_name
-                await stream.write(result)
-            except Exception as d:
-                logger.exception(f"Failed to process message:\n{d}")
-                await stream.write(TrMythicC2ToCustomMessageFormatMessageResponse(
-                    Success=False,
-                    TranslationContainerName=tr_name,
-                    Error=f"Failed to process message:\n{d}"
-                ))
+        while True:
+            stream = client.TranslateFromMythicToCustomFormat()
+            await stream.write(TrMythicC2ToCustomMessageFormatMessageResponse(
+                Success=True,
+                TranslationContainerName=tr_name
+            ))
+            logger.info(f"Connected to gRPC for handling MythicC2 to CustomC2 Translations for {tr_name}")
+            async for request in stream:
+                try:
+                    result = await translationServices[tr_name].translate_to_c2_format(request)
+                    result.TranslationContainerName = tr_name
+                    await stream.write(result)
+                except Exception as d:
+                    logger.exception(f"Failed to process message:\n{d}")
+                    await stream.write(TrMythicC2ToCustomMessageFormatMessageResponse(
+                        Success=False,
+                        TranslationContainerName=tr_name,
+                        Error=f"Failed to process message:\n{d}"
+                    ))
+            logger.error(f"disconnected from gRPC for doing mythic->custom c2 for {tr_name}")
     except Exception as e:
         logger.exception(f"[-] exception in handleMythicToCustom for {tr_name}")
diff --git a/mythic_container/__init__.py b/mythic_container/__init__.py
@@ -1,9 +1,9 @@
 from .rabbitmq import rabbitmqConnectionClass
 from .mythic_service import start_and_run_forever, test_command
 
-containerVersion = "v1.0.4"
+containerVersion = "v1.0.6"
 
-PyPi_version = "0.2.8-rc01"
+PyPi_version = "0.2.8-rc02"
 
 RabbitmqConnection = rabbitmqConnectionClass()
 

diff --git a/setup.py b/setup.py
@@ -10,7 +10,7 @@
 # This call to setup() does all the work
 setup(
     name="mythic_container",
-    version="0.2.8-rc01",
+    version="0.2.8-rc02",
     description="Functionality for Mythic Services",
     long_description=README,
     long_description_content_type="text/markdown",