chore: testing again again again

Author: grokas-splunk

src/app.py

@@ -51,7 +51,7 @@
     IMAP_SUCCESS_CONNECTIVITY_TEST,
     IMAP_VALIDATE_INTEGER_MESSAGE,
 )
-from .process_email import ProcessEmail
+from .process_email import ProcessEmail, IMAP_APP_ID
 
 logger = getLogger()
 
@@ -154,6 +154,14 @@ def __init__(self, soar: SOARClient, asset: Asset):
         self._folder_name = None
         self._is_hex = False
 
+    def debug_print(self, *args):
+        """Debug print for ProcessEmail compatibility"""
+        logger.debug(" ".join(str(arg) for arg in args))
+
+    def get_app_id(self):
+        """Return IMAP app ID for ProcessEmail compatibility"""
+        return IMAP_APP_ID
+
     def _get_error_message_from_exception(self, e):
         """Get appropriate error message from the exception"""
         error_code = None
@@ -212,19 +220,15 @@ def _connect_to_server(self, first_try=True):
         use_ssl = self.asset.use_ssl
         server = self.asset.server
 
-        # Set timeout to avoid stall
         socket.setdefaulttimeout(60)
 
-        # Connect to the server
         try:
             if is_oauth or use_ssl:
                 self._imap_conn = imaplib.IMAP4_SSL(server)
             else:
                 self._imap_conn = imaplib.IMAP4(server)
-                # Try STARTTLS for non-SSL connections
                 with contextlib.suppress(Exception):
                     self._imap_conn.starttls()
-                    logger.debug("STARTTLS enabled successfully")
         except Exception as e:
             error_text = self._get_error_message_from_exception(e)
             raise Exception(
@@ -235,25 +239,22 @@ def _connect_to_server(self, first_try=True):
 
         logger.info(IMAP_CONNECTED_TO_SERVER)
 
-        # Login
         try:
             if is_oauth:
                 auth_string = self._generate_oauth_string(
                     self.asset.username,
                     self._state.get("oauth_token", {}).get("access_token"),
                 )
-                result, data = self._imap_conn.authenticate(
+                result, _ = self._imap_conn.authenticate(
                     "XOAUTH2", lambda _: auth_string
                 )
             else:
-                result, data = self._imap_conn.login(
+                result, _ = self._imap_conn.login(
                     self.asset.username, self.asset.password
                 )
         except Exception as e:
             error_text = self._get_error_message_from_exception(e)
-            # If token is expired, use the refresh token to re-new the access token
             if first_try and is_oauth and "Invalid credentials" in error_text:
-                logger.debug("Try to generate token from refresh token")
                 self._interactive_auth_refresh()
                 return self._connect_to_server(False)
             raise Exception(
@@ -263,14 +264,12 @@ def _connect_to_server(self, first_try=True):
             ) from None
 
         if result != "OK":
-            logger.debug(f"Logging in error, result: {result} data: {data}")
             raise Exception(IMAP_ERROR_LOGGING_IN_TO_SERVER)
 
         logger.info(IMAP_LOGGED_IN)
 
-        # List imap data
         try:
-            result, data = self._imap_conn.list()
+            result, _ = self._imap_conn.list()
         except Exception as e:
             error_text = self._get_error_message_from_exception(e)
             raise Exception(
@@ -283,7 +282,7 @@ def _connect_to_server(self, first_try=True):
 
         self._folder_name = self.asset.folder
         try:
-            result, data = self._imap_conn.select(
+            result, _ = self._imap_conn.select(
                 f'"{imap_utf7.encode(self._folder_name).decode()}"', True
             )
         except Exception as e:
@@ -296,21 +295,14 @@ def _connect_to_server(self, first_try=True):
             ) from e
 
         if result != "OK":
-            logger.debug(f"Error selecting folder, result: {result} data: {data}")
             raise Exception(
                 IMAP_ERROR_SELECTING_FOLDER.format(folder=self._folder_name)
             )
 
         logger.info(IMAP_SELECTED_FOLDER.format(folder=self._folder_name))
 
-        no_of_emails = data[0]
-        logger.debug(f"Total emails: {no_of_emails}")
-
     def _get_email_data(self, muuid, folder=None, is_diff=False):
         """Get email data from IMAP server"""
-        email_data = None
-        data_time_info = None
-
         if is_diff and folder:
             try:
                 result, data = self._imap_conn.select(
@@ -325,17 +317,15 @@ def _get_email_data(self, muuid, folder=None, is_diff=False):
                 ) from e
 
             if result != "OK":
-                logger.debug(f"Error selecting folder, result: {result} data: {data}")
                 raise Exception(IMAP_ERROR_SELECTING_FOLDER.format(folder=folder))
 
             logger.info(IMAP_SELECTED_FOLDER.format(folder=folder))
 
-        # query for the whole email body
         try:
             (result, data) = self._imap_conn.uid(
                 "fetch", muuid, "(INTERNALDATE RFC822)"
             )
-        except TypeError:  # py3
+        except TypeError:
             (result, data) = self._imap_conn.uid(
                 "fetch", str(muuid), "(INTERNALDATE RFC822)"
             )
@@ -345,10 +335,6 @@ def _get_email_data(self, muuid, folder=None, is_diff=False):
                 IMAP_FETCH_ID_FAILED.format(muuid=muuid, excep=error_text)
             ) from e
 
-        logger.debug(
-            f"UID fetch result: {result}, data type: {type(data)}, data: {data}"
-        )
-
         if result != "OK":
             raise Exception(
                 IMAP_FETCH_ID_FAILED_RESULT.format(
@@ -366,6 +352,7 @@ def _get_email_data(self, muuid, folder=None, is_diff=False):
 
         if not isinstance(data[0], tuple) or len(data[0]) < 2:
             raise Exception(f"Invalid data structure for email ID {muuid}: {data[0]}")
+
         try:
             email_data = data[0][1].decode("UTF-8")
         except UnicodeDecodeError:
@@ -386,28 +373,19 @@ def _get_email_ids_to_process(self, max_emails, lower_id, manner):
             raise Exception(f"Failed to get email IDs. Server response: {data}")
 
         if not data or not data[0]:
-            logger.info("No emails found")
             return []
 
         uids = [int(uid) for uid in data[0].split()]
 
         if len(uids) == 1 and uids[0] < lower_id:
-            logger.info(f"No new UIDs found (only {uids[0]} which is < {lower_id})")
             return []
 
         uids.sort()
-
         max_emails = int(max_emails)
 
         if manner == "latest first":
-            logger.info(
-                f"Getting {max_emails} MOST RECENT email UIDs since UID {lower_id}"
-            )
-            # Return the latest (rightmost) items
             return uids[-max_emails:]
         else:
-            logger.info(f"Getting NEXT {max_emails} email UIDs since UID {lower_id}")
-            # Return the oldest (leftmost) items
             return uids[:max_emails]
 
     def _parse_and_create_artifacts(
@@ -423,9 +401,6 @@ def _parse_and_create_artifacts(
                 dt = parse_data["dt"]
                 dt.replace(tzinfo=tz.tzlocal())
                 epoch = int(time.mktime(dt.timetuple())) * 1000
-                logger.debug(f"Internal date Epoch: {dt}({epoch})")
-            else:
-                logger.debug(f"Unable to parse date/time: {data_time_info}")
 
         if config is None:
             config = {
@@ -437,35 +412,27 @@ def _parse_and_create_artifacts(
             }
 
         process_email = ProcessEmail()
+        process_email._base_connector = self
+        process_email._folder_name = self._folder_name
+        process_email._is_hex = self._is_hex
 
-        class MockConnector:
-            def __init__(self, helper):
-                self.helper = helper
-                self.containers = []
-                self.artifacts = []
-
-            def save_container(self, container_dict):
-                self.containers.append(container_dict)
-                return {"id": 0}
-
-            def save_artifact(self, artifact_dict):
-                self.artifacts.append(artifact_dict)
-                return {"id": 0}
-
-            def send_progress(self, message):
-                logger.debug(message)
-
-            def get_config(self):
-                return config
+        ret_val, message, results = process_email._int_process_email(
+            email_data, email_id, epoch
+        )
 
-        mock = MockConnector(self)
-        process_email.process_email(mock, email_data, email_id, config, epoch)
+        if not ret_val:
+            logger.error(f"Failed to process email {email_id}: {message}")
+            return
 
-        for container_dict in mock.containers:
-            yield Container(**container_dict)
+        for result in results:
+            container_dict = result.get("container")
+            if container_dict:
+                yield Container(**container_dict)
 
-        for artifact_dict in mock.artifacts:
-            yield Artifact(**artifact_dict)
+            artifacts = result.get("artifacts", [])
+            for artifact_dict in artifacts:
+                if artifact_dict:
+                    yield Artifact(**artifact_dict)
 
 
 @app.test_connectivity()
@@ -488,70 +455,52 @@ def on_poll(
     params: OnPollParams, soar: SOARClient, asset: Asset
 ) -> Iterator[Container | Artifact]:
     """Poll for new emails and ingest as containers/artifacts"""
-
     helper = ImapHelper(soar, asset)
-
-    # Connect to server
     helper._connect_to_server()
 
-    # Access ingestion state for tracking email processing
     state = app.actions_manager.ingestion_state
 
-    # Determine if this is first run
-    is_first_run = state.get("first_run", True)
-    lower_id = state.get("next_muid", 1)
+    is_poll_now = params.container_count is not None
 
-    # Determine max emails based on first run
-    if is_first_run:
-        max_emails = int(asset.first_run_max_emails)
+    if is_poll_now:
+        lower_id = 1
+        max_emails = params.container_count if params.container_count > 0 else 100
     else:
-        max_emails = int(asset.max_emails)
-
-    # Allow container_count to override for manual polling
-    if params.container_count:
-        max_emails = params.container_count
-
-    logger.info(f"Will fetch up to {max_emails} emails (starting from UID {lower_id})")
-
-    # Get email IDs to process
-    try:
-        email_ids = helper._get_email_ids_to_process(
-            max_emails, lower_id, asset.ingest_manner
+        is_first_run = state.get("first_run", True)
+        lower_id = state.get("next_muid", 1)
+        max_emails = (
+            int(asset.first_run_max_emails) if is_first_run else int(asset.max_emails)
         )
-        logger.info(f"Got {len(email_ids)} email IDs: {email_ids}")
-    except Exception as e:
-        logger.error(f"Failed to get email IDs: {e}")
-        logger.exception("Full traceback:")
-        raise
+
+    email_ids = helper._get_email_ids_to_process(
+        max_emails, lower_id, asset.ingest_manner
+    )
 
     if not email_ids:
         logger.info("No new emails to ingest")
         return
 
-    logger.info(f"Found {len(email_ids)} emails to ingest")
-
-    for i, email_id in enumerate(email_ids):
-        logger.info(f"Processing email UID {email_id} ({i + 1}/{len(email_ids)})")
-
+    for email_id in email_ids:
         try:
             email_data, data_time_info = helper._get_email_data(email_id)
-            logger.info(f"Got email data, size: {len(email_data)} bytes")
 
             yield from helper._parse_and_create_artifacts(
                 email_id, email_data, data_time_info, asset
             )
 
         except Exception as e:
             logger.error(f"Error processing email {email_id}: {e}")
-            logger.exception("Full traceback:")
             continue
 
-    if email_ids:
+    if email_ids and not is_poll_now:
         state["next_muid"] = int(email_ids[-1]) + 1
         state["first_run"] = False
-        logger.info(f"Updated next_muid to {state['next_muid']}")
 
-    logger.info(f"On_poll completed, processed {len(email_ids)} emails")
+
+class GetEmailSummary(ActionOutput):
+    """Summary output for get_email action"""
+
+    container_id: int | None = None
 
 
 class GetEmailParams(Params):
@@ -778,9 +727,12 @@ def get_email(params: GetEmailParams, soar: SOARClient, asset: Asset) -> GetEmai
                         soar._artifacts_api.create(artifact_dict)
 
             message = f"Email ingested with container ID: {container_id}"
+            soar.set_summary(GetEmailSummary(container_id=container_id))
         else:
             message = "Email not ingested."
 
+        soar.set_message(message)
+
         ret_val = {"message": message}
         if container_id:
             ret_val["container_id"] = container_id

src/process_email.py

@@ -28,7 +28,17 @@
 
 from bs4 import BeautifulSoup, UnicodeDammit
 from requests.structures import CaseInsensitiveDict
-from soar_sdk.shims import phantom
+
+try:
+    import phantom.app as phantom
+except ImportError:
+    from soar_sdk.shims import phantom
+    from soar_sdk.shims.phantom.app import APP_SUCCESS, APP_ERROR
+
+    phantom.APP_SUCCESS = APP_SUCCESS
+    phantom.APP_ERROR = APP_ERROR
+    phantom.is_fail = lambda x: x == APP_ERROR
+    phantom.is_success = lambda x: x == APP_SUCCESS
 
 from . import phantom_rules
 from . import phantom_utils as ph_utils