From 1e46a0863a3dd7296b44f187d3a70894e10e9d5c Mon Sep 17 00:00:00 2001
From: mr-tz <moritz.raabe@mandiant.com>
Date: Mon, 8 Apr 2024 17:25:51 +0200
Subject: [PATCH] check runtime using Python instead of Bash

---
 .github/check_runtimes.py   | 64 +++++++++++++++++++++++++++++++++++++
 .github/workflows/tests.yml | 22 +------------
 2 files changed, 65 insertions(+), 21 deletions(-)
 create mode 100644 .github/check_runtimes.py

diff --git a/.github/check_runtimes.py b/.github/check_runtimes.py
new file mode 100644
index 0000000..0db7f93
--- /dev/null
+++ b/.github/check_runtimes.py
@@ -0,0 +1,64 @@
+# Copyright (C) 2024 Mandiant, Inc. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at: [package root]/LICENSE.txt
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+"""
+Check runtime of testfiles.
+"""
+
+import sys
+import time
+import logging
+import argparse
+from pathlib import Path
+
+import capa.main
+
+logger = logging.getLogger("capa.tests.data")
+
+THRESHOLD = 60 * 3
+TARGET_EXTS = (".exe_", ".dll_", ".elf_", ".sys_", ".raw32", ".raw64", ".BinExport")
+IGNORED_DIRS = ("aarch64",)
+
+
+def main(argv=None):
+    if argv is None:
+        argv = sys.argv[1:]
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("files", nargs="+", help="Paths of added/modified files")
+    args = parser.parse_args(args=argv)
+
+    test_failed = False
+    for file in args.files:
+        file = Path(file)
+        # Skip ignored directories
+        if any((ignored_dir in file.parts) for ignored_dir in IGNORED_DIRS):
+            continue
+
+        if not file.name.endswith(TARGET_EXTS):
+            continue
+
+        time0 = time.time()
+        capa.main.main(["-q", "-v", str(file)])
+        diff = time.time() - time0
+
+        if diff > THRESHOLD:
+            logger.info("capa ran for %s seconds, please provide a different sample so we can test more quickly", diff)
+            test_failed = True
+        else:
+            logger.info("all good, capa ran for %s seconds", diff)
+
+    if test_failed:
+        return 1
+    else:
+        logger.info("test files look good!")
+        return 0
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    sys.exit(main())
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index f4edd04..b48faf7 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -49,24 +49,4 @@ jobs:
       with:
         format: 'csv'
     - name: Check capa runtime on modified files
-      run: |
-        THRESHOLD=180
-        EXCLUSION_REGEX="aarch64"
-        FILE_EXTENSION_REGEX=".exe_|.dll_|.elf_|.sys_|.raw32|.raw64"
-        exitcode=0
-        cd tests/data
-        mapfile -d ',' -t added_modified_files < <(printf '%s,' '${{ steps.files.outputs.all }}')
-        for changed_file in "${added_modified_files[@]}"; do
-          if [[ $changed_file =~ $FILE_EXTENSION_REGEX && ! $changed_file =~ $EXCLUSION_REGEX ]]; then
-            time0=$SECONDS
-            capa -q -v "$changed_file"
-            diff=$(($SECONDS-time0))
-            if [[ $diff -gt $THRESHOLD ]]; then
-              echo "capa ran for $diff seconds, please provide a different sample so we can test more quickly"
-              exitcode=1
-            else
-              echo "all good, capa ran for $diff seconds"
-            fi
-          fi
-        done
-        exit $exitcode
+      run: python .github/check_runtimes.py ${{ steps.files.outputs.all }}