From adf0c27d14340b61e6b5fa25cff7665d1c35c36d Mon Sep 17 00:00:00 2001
From: Christian Herdtweck <christian.herdtweck@intra2net.com>
Date: Mon, 19 Jun 2023 11:40:50 +0200
Subject: [PATCH] Move mail util for waiting for quarantine from QA

---
 src/arnied_wrapper.py |   46 ++++++++++++++++++++++++++++++++++++++++++++--
 src/mail_utils.py     |    2 +-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/arnied_wrapper.py b/src/arnied_wrapper.py
index 5387a15..3830843 100644
--- a/src/arnied_wrapper.py
+++ b/src/arnied_wrapper.py
@@ -30,12 +30,12 @@ Copyright: Intra2net AG
 """
 
 import os
-import sys
 import time
-import re
 import subprocess
 import shutil
 import tempfile
+from shlex import quote
+from typing import Any
 import logging
 log = logging.getLogger('pyi2ncommon.arnied_wrapper')
 
@@ -247,6 +247,48 @@ def wait_for_email_transfer(timeout=300, vm=None):
                 .format(timeout))
 
 
+def wait_for_quarantine_processing(vm_session: Any = None, max_wait: int = 30) -> bool:
+    """
+    Wait until quarantined is finished processing.
+
+    This checks quarantined's input and temp dirs and returns as soon as they are all empty or when
+    max waiting time is reached.
+
+    To be used after :py:func:`wait_for_email_transfer`.
+
+    :param vm_session: optional :py:class:`aexpect.client.ShellSession`; default: run on localhost
+    :param max_wait: maximum time in seconds to wait here
+    :returns: `True` if all quarantines have empty input/tmp dirs upon return, `False` if we
+              reached `max_time` while waiting
+    """
+    def has_files(dirname: str) -> bool:
+        # Quick abstraction to check dir on local host or in remote session
+        if vm_session is None:
+            return bool(os.listdir(dirname))
+        cmd = f"ls -UNq {quote(dirname)}"
+        status, output = vm_session.cmd_status_output(cmd)
+        if status == 0:
+            return bool(output.strip())    # False <==> empty output <==> no files
+        elif status == 2:   # dir does not exist
+            return False    # non-existent dir is empty
+        else:
+            raise RuntimeError(f"{cmd} returned {status} and output: {output}")
+
+    n_sleep = 0
+    for quarantine in ("spam", "attachment", "virus"):
+        for subdir in ("q-in", "q-tmp"):
+            try:
+                full_dir = f"/datastore/quarantine/{quarantine}/{subdir}/"
+                while has_files(full_dir):
+                    n_sleep += 1
+                    if n_sleep > max_wait:
+                        return False     # abort
+                    time.sleep(1)
+            except FileNotFoundError:    # no such directory on local host
+                continue
+    return True
+
+
 def schedule(program, exec_time=0, optional_args="", vm=None):
     """
     Schedule a program to be executed at a given unix time stamp.
diff --git a/src/mail_utils.py b/src/mail_utils.py
index 5545546..2ad3e25 100644
--- a/src/mail_utils.py
+++ b/src/mail_utils.py
@@ -34,7 +34,7 @@ from email.utils import parsedate_to_datetime
 from email.parser import BytesParser
 from email import policy
 
-# outsourced source, import required for compatiblity
+# outsourced source, import required for compatibility
 from .imap_mailbox import ImapMailbox           # pylint: disable=unused-import
 from .mail_validator import *                   # pylint: disable=unused-import
 from .sysmisc import replace_file_regex
-- 
1.7.1