Add do-not-auto-tag flag CLI option (#28)

* Add `do-not-auto-tag` CLI option flag to manage `skip-scraping` on APIs
* Fix path issues in test env

.github/workflows/python-package.yml

@@ -40,6 +40,7 @@ jobs:
     - name: Test with pytest
       env: # Set the secret as an environment variable
         REAI_API_KEY: ${{ secrets.REAI_API_KEY }}
+        REAI_API_HOST: ${{ secrets.REAI_API_HOST }}
       run: |
         pip install pytest
         python -m pytest tests/

src/reait/api.py

@@ -10,7 +10,7 @@
 from datetime import datetime
 
 from sklearn.metrics.pairwise import cosine_similarity
-from os import access, R_OK
+from os import access, R_OK, environ
 from os.path import basename, isfile, expanduser, getsize
 from requests import request, Response, HTTPError
 from numpy import array, vstack, dot, arccos, pi
@@ -20,9 +20,9 @@
 __version__ = "1.0.1"
 
 re_conf = {
-    "apikey": "l1br3",
-    "host": "https://api.reveng.ai",
-    "model": "binnet-0.3-x86",
+    "apikey": environ.get("REAI_API_KEY", "l1br3"),
+    "host": environ.get("REAI_API_HOST", "https://api.reveng.ai"),
+    "header-host": environ.get("HEADER_HOST", None),
 }
 
 
@@ -41,11 +41,11 @@ def __init__(self, reason: str, end_point: str = None):
         super().__init__(reason, response=response)
 
 
-def reveng_req(r: request, end_point: str, data: dict = None, ex_headers: dict = None,
+def reveng_req(req: request, end_point: str, data: dict = None, ex_headers: dict = None,
                params: dict = None, json_data: dict = None, timeout: int = 60, files: dict = None) -> Response:
     """
     Constructs and sends a Request
-    :param r: Method for the new Request
+    :param req: Method for the new Request
     :param end_point: Endpoint to add to the base URL
     :param ex_headers: Extended HTTP headers to add
     :param data: Dictionary, list of tuples, bytes, or file-like object to send in the body
@@ -56,17 +56,19 @@ def reveng_req(r: request, end_point: str, data: dict = None, ex_headers: dict =
     """
     url = f"{re_conf['host']}/{end_point if end_point[0] != '/' else end_point[1:]}"
     headers = {"Authorization": re_conf["apikey"]}
+    if re_conf["header-host"] is not None:
+        headers["Host"] = re_conf["header-host"]
 
     if ex_headers:
         headers.update(ex_headers)
 
     logger.debug("Making %s request %s:\n  - headers: %s\n  - data: %s\n  - json_data: %s\n  - params: %s\n  - files: %s",
-                 r.__name__.upper(), url, headers, data, json_data, params, files)
+                 req.__name__.upper(), url, headers, data, json_data, params, files)
 
-    response: Response = r(url, headers=headers, json=json_data, data=data, params=params, timeout=timeout, files=files)
+    response: Response = req(url, headers=headers, json=json_data, data=data, params=params, timeout=timeout, files=files)
 
     logger.debug("Making %s response %s:\n  - headers: %s\n  - status_code: %d\n  - content: %s",
-                 r.__name__.upper(), url, response.headers, response.status_code, response.text)
+                 req.__name__.upper(), url, response.headers, response.status_code, response.text)
 
     return response
 
@@ -166,7 +168,8 @@ def RE_delete(fpath: str, binary_id: int = 0) -> Response:
 def RE_analyse(fpath: str, model_name: str = None, isa_options: str = None,
                platform_options: str = None, file_options: str = None, dynamic_execution: bool = False,
                command_line_args: str = None, binary_scope: str = None, tags: list = None, priority: int = 0,
-               duplicate: bool = False, symbols: dict = None, debug_fpath: str = None) -> Response:
+               duplicate: bool = False, symbols: dict = None, debug_fpath: str = None,
+               skip_scraping: bool = False) -> Response:
     """
     Start analysis job for binary file
     :param fpath: File path for binary to analyse
@@ -182,6 +185,7 @@ def RE_analyse(fpath: str, model_name: str = None, isa_options: str = None,
     :param duplicate: Duplicate an existing binary
     :param symbols: JSON object containing the base address and the list of functions
     :param debug_fpath: File path for debug file
+    :param skip_scraping: Disable/Enable auto-tagging of binary sample in relevant APIs
     """
     bin_id = re_binary_id(fpath)
     result = re_hash_check(bin_id)
@@ -207,7 +211,8 @@ def RE_analyse(fpath: str, model_name: str = None, isa_options: str = None,
             pass
 
     for p_name in ("model_name", "isa_options", "platform_options", "file_options",
-                   "dynamic_execution", "command_line_args", "binary_scope", "tags", "priority", "symbols",):
+                   "dynamic_execution", "command_line_args", "binary_scope",
+                   "tags", "priority", "symbols", "skip_scraping"):
         p_value = locals()[p_name]
 
         if p_value:

src/reait/main.py

@@ -17,7 +17,7 @@
 from glob import iglob
 import numpy as np
 
-import api
+from . import api
 
 rerr = Console(file=stderr, width=180)
 rout = Console(file=stdout, width=180)
@@ -210,6 +210,7 @@ def main() -> int:
                         help="Override analysis visibility (scope). Valid values are 'public' or 'private'[DEFAULT]")
     parser.add_argument("--tags", default=None, type=str,
                         help="Assign tags to an analysis. Valid responses are tag1,tag2,tag3.")
+    parser.add_argument("--do-not-auto-tag", default=False, action="store_true" ,help="Disable auto-tagging in API views",)
     parser.add_argument("--priority", default=0, type=int, help="Add priority to processing queue.")
     parser.add_argument("--verbose", default=False, action="store_true", help="Set verbose output.")
     parser.add_argument("--debug", default=None, help="Debug file path to write pass with analysis")
@@ -266,7 +267,9 @@ def main() -> int:
                                platform_options=args.platform, dynamic_execution=args.dynamic_execution,
                                command_line_args=args.cmd_line_args, file_options=args.exec_format,
                                binary_scope=args.scope.upper(), tags=tags, priority=args.priority,
-                               duplicate=args.duplicate, debug_fpath=args.debug)
+                               duplicate=args.duplicate, debug_fpath=args.debug,
+                               skip_scraping=args.do_not_auto_tag
+                               )
 
             if args.delete:
                 try:
@@ -306,7 +309,9 @@ def main() -> int:
                            platform_options=args.platform, dynamic_execution=args.dynamic_execution,
                            command_line_args=args.cmd_line_args, file_options=args.exec_format,
                            binary_scope=args.scope.upper(), tags=tags, priority=args.priority,
-                           duplicate=args.duplicate, debug_fpath=args.debug)
+                           duplicate=args.duplicate, debug_fpath=args.debug,
+                           skip_scraping=args.do_not_auto_tag
+                           )
 
         elif args.extract:
             embeddings = api.RE_embeddings(args.binary).json()

tests/utils.py

@@ -28,10 +28,12 @@
 class BaseTestCase(TestCase):
     __metaclass__ = abc.ABCMeta
 
+    MODEL_NAME_PREFIX = "binnet-0.4-x86-"
     _cwd: str = None
     _fpath: str = None
     _platform: str = None
 
+
     @classmethod
     def setUpClass(cls):
         cls._cwd = join(CWD, "binaries")
@@ -43,10 +45,10 @@ def setUpClass(cls):
 
         testlog.info("Random selection of '%s' in binaries folder", cls._fpath)
 
-        api.re_conf["model"] = "binnet-0.3-x86-" + cls._platform
-
+        api.re_conf["model"] = f"{cls.MODEL_NAME_PREFIX}{cls._platform}"
         # Get the API key from the environment variable
         api.re_conf["apikey"] = getenv("REAI_API_KEY", api.re_conf["apikey"])
+        api.re_conf["header-host"] = getenv("HEADER_HOST", 'api.local')
 
         # Deletes all previous analyses from the RevEng.AI platform
         cls._cleanup_binaries(cls._fpath)
@@ -58,7 +60,12 @@ def _cleanup_binaries(fpath: str) -> None:
 
             testlog.info("Getting all previous analyses for %s", bin_id)
 
-            response = api.reveng_req(get, "v1/search", json_data={"sha_256_hash": bin_id}).json()
+            response = api.reveng_req(
+                get,
+                "v1/search",
+                json_data={"sha_256_hash": bin_id},
+                ex_headers={"Host": "api.local"}
+            ).json()
 
             if not response["success"]:
                 testlog.error("Failed to get all previous analysis.\n%s", response)