reait 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

reait/api.py

@@ -10,12 +10,12 @@ import tomli
 from datetime import datetime
 from hashlib import sha256
 from lief import parse, Binary, ELF, PE, MachO
-from numpy import array, vstack, dot, arccos, pi
+from numpy import array, vstack
 from pandas import DataFrame
 from requests import request, Response, HTTPError
 from sklearn.metrics.pairwise import cosine_similarity
 
-__version__ = "1.1.0"
+__version__ = "1.2.0"
 
 re_conf = {
     "apikey": environ.get("REAI_API_KEY", ""),
@@ -32,7 +32,8 @@ class ReaitError(HTTPError):
 
         response.reason = reason
         response.status_code = 404
-        response._content = b'{"success": false, "error": "' + reason.encode() + b'"}'
+        response._content = b'{"success": false, "error": "' + \
+            reason.encode() + b'"}'
         response.url = (
             f"{re_conf['host']}/{end_point if end_point[0] != '/' else end_point[1:]}"
             if end_point
@@ -243,6 +244,9 @@ def RE_analyse(
     symbols: dict = None,
     debug_fpath: str = None,
     skip_scraping: bool = False,
+    skip_capabilities: bool = False,
+    skip_sbom: bool = False,
+    advanced_analysis: bool = False
 ) -> Response:
     """
     Start analysis job for binary file
@@ -302,6 +306,9 @@ def RE_analyse(
         "priority",
         "symbols",
         "skip_scraping",
+        "skip_capabilities",
+        "skip_sbom",
+        "advanced_analysis"
     ):
         p_value = locals()[p_name]
 
@@ -315,7 +322,8 @@ def RE_analyse(
         )
     elif res.status_code == 400:
         if "error" in res.json().keys():
-            logger.warning("Error analysing %s - %s", fpath, res.json()["error"])
+            logger.warning("Error analysing %s - %s",
+                           fpath, res.json()["error"])
 
     res.raise_for_status()
     return res
@@ -330,7 +338,8 @@ def RE_upload(fpath: str) -> Response:
     result = re_hash_check(bin_id)
 
     if result:
-        logger.info("File %s - %s already uploaded. Skipping upload...", fpath, bin_id)
+        logger.info(
+            "File %s - %s already uploaded. Skipping upload...", fpath, bin_id)
 
         res = Response()
         res.status_code = 200
@@ -346,7 +355,8 @@ def RE_upload(fpath: str) -> Response:
         )
     else:
         with open(fpath, "rb") as fd:
-            res: Response = reveng_req(requests.post, "v1/upload", files={"file": fd})
+            res: Response = reveng_req(
+                requests.post, "v1/upload", files={"file": fd})
 
         if res.ok:
             logger.info(
@@ -354,7 +364,8 @@ def RE_upload(fpath: str) -> Response:
             )
         elif res.status_code == 400:
             if "error" in res.json().keys():
-                logger.warning("Error uploading %s - %s", fpath, res.json()["error"])
+                logger.warning("Error uploading %s - %s",
+                               fpath, res.json()["error"])
         elif res.status_code == 413:
             logger.warning("File too large. Please upload files under 10MB.")
         elif res.status_code == 500:
@@ -511,7 +522,8 @@ def RE_compute_distance(embedding: list, embeddings: list, nns: int = 5) -> list
 def RE_nearest_symbols_batch(
     function_ids: list[int],
     nns: int = 5,
-    collections: list[str] = None,
+    collections: list[int] = None,
+    binaries: list[int] = None,
     distance: float = 0.1,
     debug_enabled: bool = False,
 ) -> Response:
@@ -531,10 +543,13 @@ def RE_nearest_symbols_batch(
     }
 
     if collections:
-        # api param is collection, not collections
-        params["collection"] = collections
+        params["collection_search_list"] = collections
+
+    if binaries:
+        params["binaries_search_list"] = binaries
 
-    res: Response = reveng_req(requests.post, "v1/ann/symbol/batch", json_data=params)
+    res: Response = reveng_req(
+        requests.post, "v1/ann/symbol/batch", json_data=params)
 
     res.raise_for_status()
     return res
@@ -655,9 +670,11 @@ def RE_functions_rename(function_id: int, new_name: str) -> Response:
     )
 
     if res.ok:
-        logger.info("FunctionId %d has been renamed with '%s'.", function_id, new_name)
+        logger.info("FunctionId %d has been renamed with '%s'.",
+                    function_id, new_name)
     else:
-        logger.warning("Error, cannot rename FunctionId %d. %s", function_id, res.text)
+        logger.warning("Error, cannot rename FunctionId %d. %s",
+                       function_id, res.text)
 
     res.raise_for_status()
     return res
@@ -742,7 +759,9 @@ def RE_functions_list(
         params["max_v_address"] = max_v_address
 
     res: Response = reveng_req(
-        requests.get, f"v2/analyses/{analysis_id}/info/functions/list", params=params
+        requests.get,
+        f"v2/analyses/{analysis_id}/info/functions/list",
+        params=params
     )
 
     res.raise_for_status()
@@ -755,7 +774,8 @@ def RE_function_callers_callees(function: int) -> Response:
     Get the callers and callees of a functions
     :param function: Function ID
     """
-    res: Response = reveng_req(requests.get, f"v2/functions/{function}/callees_callers")
+    res: Response = reveng_req(
+        requests.get, f"v2/functions/{function}/callees_callers")
 
     res.raise_for_status()
     return res
@@ -766,7 +786,8 @@ def RE_analysis_info(analysis_id: int) -> Response:
     Get the analysis information
     :param analysis_id: Analysis ID
     """
-    res: Response = reveng_req(requests.get, f"v2/analyses/{analysis_id}/info/basic")
+    res: Response = reveng_req(
+        requests.get, f"v2/analyses/{analysis_id}/info/basic")
 
     res.raise_for_status()
     return res
@@ -847,10 +868,12 @@ def _binary_format(binary: Binary) -> str:
         return "Mach-O"
 
     logger.error(
-        "Error, could not determine or unsupported" f" binary format: {binary.format}."
+        "Error, could not determine or unsupported"
+        f" binary format: {binary.format}."
     )
     raise RuntimeError(
-        "Error, could not determine or " f"unsupported binary format: {binary.format}"
+        "Error, could not determine or "
+        f"unsupported binary format: {binary.format}"
     )
 
 
@@ -914,12 +937,47 @@ def RE_analysis_id(fpath: str, binary_id: int = 0) -> Response:
     return res
 
 
-def RE_generate_data_types(analysis_id: int, function_ids: list[int]) -> Response:
+def RE_functions_data_types(
+    function_ids: list[int],
+) -> Response:
+    """
+    Get data types for the functions
+    :param functions_ids: List of function IDs
+    :return: Response object
+    """
+    endpoint = "/v2/functions/data_types"
+    res: Response = reveng_req(
+        requests.post, endpoint, json_data={"function_ids": function_ids}
+    )
+    res.raise_for_status()
+    return res
+
+
+def RE_functions_data_types_poll(
+    function_ids: list[int],
+) -> Response:
+    """
+    Poll data types for the functions
+    :param functions_ids: List of function IDs
+    :return: Response object
+    """
+    endpoint = "/v2/functions/data_types"
+    res: Response = reveng_req(
+        requests.get, endpoint, params={"function_ids": function_ids}
+    )
+    res.raise_for_status()
+    return res
+
+
+def RE_generate_data_types(
+    analysis_id: int,
+    function_ids: list[int]
+) -> Response:
     """
     Generate data types for the analysis
     :param aid: Analysis ID
     """
-    end_point = f"/v2/analyses/{analysis_id}/info/functions/data_types"
+    end_point = f"/v2/analyses/{analysis_id}/functions/data_types"
 
     res: Response = reveng_req(
         requests.post, end_point, json_data={"function_ids": function_ids}
@@ -928,13 +986,28 @@ def RE_generate_data_types(analysis_id: int, function_ids: list[int]) -> Respons
     return res
 
 
+def RE_poll_data_types(
+    analysis_id: int,
+    function_id: int,
+) -> Response:
+    """
+    Poll data types for the analysis
+    :param aid: Analysis ID
+    """
+    end_point = f"/v2/analyses/{analysis_id}/functions/{function_id}/data_types"
+
+    res: Response = reveng_req(requests.get, end_point)
+    res.raise_for_status()
+    return res
+
+
 def RE_list_data_types(analysis_id: int, function_ids: list[int]) -> Response:
     """
     List data types for the analysis
     :param aid: Analysis ID
     :param function_ids: List of function IDs
     """
-    end_point = f"/v2/analyses/{analysis_id}/info/functions/data_types"
+    end_point = f"/v2/analyses/{analysis_id}/functions/data_types"
 
     res: Response = reveng_req(
         requests.get, end_point, json_data={"function_ids": function_ids}
@@ -959,16 +1032,25 @@ def RE_begin_ai_decompilation(function_id: int) -> Response:
     return res
 
 
-def RE_poll_ai_decompilation(function_id: int) -> Response:
+def RE_poll_ai_decompilation(
+    function_id: int,
+    summarise: bool = False
+) -> Response:
     """
     Poll AI decompilation for the function
     :param function_id: Function ID
     """
     end_point = f"/v2/functions/{function_id}/ai-decompilation"
 
+    params = {}
+
+    if summarise:
+        params["summarise"] = summarise
+
     res: Response = reveng_req(
         requests.get,
         end_point,
+        params=params,
     )
     res.raise_for_status()
     return res
@@ -988,20 +1070,80 @@ def RE_analysis_lookup(binary_id: int) -> Response:
 def RE_collections_search(
         page: int = 1,
         page_size: int = 10,
-        partial_collection_name: str = "",
-        partial_binary_name: str = "",
-        partial_binary_sha256: str = "",
-        tags: list[str] | str = "",
-        model_name: str = "",
+        query: dict = {},
 ) -> Response:
     """
+    Search for collections in the database
     """
     end_point = "/v2/search/collections"
-    res: Response = reveng_req(requests.get, end_point, params={
+    params = {
         "page": page,
         "page_size": page_size,
-        "partial_collection_name": partial_collection_name,
-    })
+    }
+
+    # this api support:
+    #   partial_collection_name
+    #   partial_binary_name
+    #   partial_binary_sha256
+    #   model_name
+    #   tags
+
+    def exist_and_populated(key: str) -> bool:
+        return key in query and query[key] is not None and len(query[key]) > 0
+
+    if exist_and_populated("collection_name"):
+        params["partial_collection_name"] = query["collection_name"]
+    elif exist_and_populated("binary_name"):
+        params["partial_binary_name"] = query["binary_name"]
+    elif exist_and_populated("sha_256_hash"):
+        params["partial_binary_sha256"] = query["sha_256_hash"]
+    elif exist_and_populated("tags"):
+        params["tags"] = query["tags"]
+    elif exist_and_populated("model_name"):
+        params["model_name"] = query["model_name"]
+    elif exist_and_populated("query"):
+        params["partial_collection_name"] = query["query"]
+
+    res: Response = reveng_req(requests.get, end_point, params=params)
+    res.raise_for_status()
+    return res
+
+
+def RE_binaries_search(
+        page: int = 1,
+        page_size: int = 10,
+        query: dict = {},
+) -> Response:
+    """
+    Search for binaries in the database
+    """
+    end_point = "/v2/search/binaries"
+    params = {
+        "page": page,
+        "page_size": page_size,
+    }
+
+    # this api support:
+    #   partial_name
+    #   partial_sha256
+    #   tags
+    #   model_name
+
+    def exist_and_populated(key: str) -> bool:
+        return key in query and query[key] is not None and len(query[key]) > 0
+
+    if exist_and_populated("binary_name"):
+        params["partial_name"] = query["binary_name"]
+    elif exist_and_populated("sha_256_hash"):
+        params["partial_sha256"] = query["sha_256_hash"]
+    elif exist_and_populated("tags"):
+        params["tags"] = query["tags"]
+    elif exist_and_populated("model_name"):
+        params["model_name"] = query["model_name"]
+    elif exist_and_populated("query"):
+        params["partial_name"] = query["query"]
+
+    res: Response = reveng_req(requests.get, end_point, params=params)
     res.raise_for_status()
     return res
 

{reait-1.1.0.dist-info → reait-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reait
-Version: 1.1.0
+Version: 1.2.0
 Summary: RevEng.AI Toolkit and Python API
 Home-page: https://github.com/RevEng-AI/reait
 Author: James Patrick-Evans

reait-1.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+reait/__init__.py,sha256=EoVCKwQwWxEBfwe-iEE5rFKvhi1gPEA8NPhnzXJTb2Y,42
+reait/api.py,sha256=OzRPRgXAfOUmdPWx7b7_CwKfZ6OoT1VK07VYTSmWnB4,36646
+reait/main.py,sha256=rSzEowDrK2KFmmLdbRNVsfVpvMLZNXA3fQOrBw03T4Y,20396
+reait-1.2.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reait-1.2.0.dist-info/METADATA,sha256=AQusdz1uqsRJ1qdlnlV10OlEKGlaR0n-L5vk1pt8_u4,47663
+reait-1.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+reait-1.2.0.dist-info/entry_points.txt,sha256=8Ek11o7a6O8hjBFw6-1vmkBfbv_45O2vOKj5CDUB1e4,42
+reait-1.2.0.dist-info/top_level.txt,sha256=EnJssmctKe3Ugjcmu66L9_Q4elLdAwaXK6M8E6E8f_M,6
+reait-1.2.0.dist-info/RECORD,,

reait-1.1.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-reait/__init__.py,sha256=EoVCKwQwWxEBfwe-iEE5rFKvhi1gPEA8NPhnzXJTb2Y,42
-reait/api.py,sha256=gYF4K9fns47bkIUu9_ZV1glRBqEP-J1OaIn99prRugE,33128
-reait/main.py,sha256=rSzEowDrK2KFmmLdbRNVsfVpvMLZNXA3fQOrBw03T4Y,20396
-reait-1.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reait-1.1.0.dist-info/METADATA,sha256=ULUT4H4-BSwIOVl9_cn51PwU7D_2U7qnlgT6nWAEjG0,47663
-reait-1.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-reait-1.1.0.dist-info/entry_points.txt,sha256=8Ek11o7a6O8hjBFw6-1vmkBfbv_45O2vOKj5CDUB1e4,42
-reait-1.1.0.dist-info/top_level.txt,sha256=EnJssmctKe3Ugjcmu66L9_Q4elLdAwaXK6M8E6E8f_M,6
-reait-1.1.0.dist-info/RECORD,,