From 3f525149c540c1610414cecba891f980b0b286c3 Mon Sep 17 00:00:00 2001
From: Haseeb Rabbani <haseeb.rabbani@gmail.com>
Date: Thu, 12 Oct 2023 19:16:03 -0400
Subject: [PATCH] python external bot support

---
 python-sdk/src/forta_agent/__init__.py   |  2 +-
 python-sdk/src/forta_agent/alerts_api.py | 70 ++++++++++++++++++++++++
 python-sdk/src/forta_agent/finding.py    |  4 +-
 python-sdk/src/forta_agent/utils.py      |  2 +
 4 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/python-sdk/src/forta_agent/__init__.py b/python-sdk/src/forta_agent/__init__.py
index 2028629..532edd2 100644
--- a/python-sdk/src/forta_agent/__init__.py
+++ b/python-sdk/src/forta_agent/__init__.py
@@ -11,7 +11,7 @@
 from .network import Network
 from .bloom_filter import BloomFilter
 from .utils import get_json_rpc_url, create_block_event, create_transaction_event, create_alert_event, get_web3_provider, keccak256, get_transaction_receipt, get_chain_id, get_bot_owner, get_bot_id
-from .alerts_api import get_alerts
+from .alerts_api import get_alerts, send_alerts
 from .labels_api import get_labels
 from .jwt import fetch_jwt, decode_jwt, verify_jwt, MOCK_JWT
 from web3 import Web3
diff --git a/python-sdk/src/forta_agent/alerts_api.py b/python-sdk/src/forta_agent/alerts_api.py
index e61e3ae..2846a59 100644
--- a/python-sdk/src/forta_agent/alerts_api.py
+++ b/python-sdk/src/forta_agent/alerts_api.py
@@ -1,8 +1,78 @@
 import requests
+import json
 from .alert import Alert
+from .finding import FindingType, FindingSeverity
+from .label import EntityType
 from .utils import get_forta_api_headers, get_forta_api_url
 
 
+def send_alerts(alerts):
+    if not isinstance(alerts, list):
+        alerts = [alerts]
+
+    alerts_api_url = get_forta_api_url()
+    headers = get_forta_api_headers()
+    mutation = SendAlertsRequest(alerts).get_mutation()
+
+    response = requests.request(
+        "POST", alerts_api_url, json=mutation, headers=headers)
+
+    if response.status_code == 200:
+        data = response.json().get('data')
+        if data:
+            return data.get('sendAlerts').get('alerts')
+    else:
+        message = response.text
+        raise Exception(message)
+
+
+class SendAlertsRequest:
+    def __init__(self, alerts):
+        self.alerts = []
+        # serialize the alerts list
+        for alert in alerts:
+            # convert finding timestamp to RFC3339 format
+            alert["finding"].timestamp = alert["finding"].timestamp.astimezone(
+            ).isoformat()
+            # serialize finding
+            finding = json.loads(alert["finding"].toJson())
+            # convert enums to all caps to match graphql enums
+            finding["type"] = FindingType(finding["type"]).name.upper()
+            finding["severity"] = FindingSeverity(
+                finding["severity"]).name.upper()
+            for label in finding.get("labels", []):
+                label["entityType"] = EntityType(
+                    label["entityType"]).name.upper()
+            # remove protocol field (not part of graphql schema)
+            del finding["protocol"]
+            # remove any empty-value or snake-case-keyed fields
+            finding = {k: v for k, v in finding.items()
+                       if v is not None and "_" not in k}
+            for index, label in enumerate(finding.get("labels", [])):
+                finding["labels"][index] = {k: v for k, v in label.items()
+                                            if v is not None and "_" not in k}
+            self.alerts.append({
+                "botId": alert["bot_id"],
+                "finding": finding
+            })
+
+    def get_mutation(self):
+        mutation = """
+        mutation SendAlerts($alerts: [AlertRequestInput!]!) {
+            sendAlerts(alerts: $alerts) {
+                alerts {
+                    alertHash
+                    error {
+                        code
+                        message
+                    }
+                }
+            }
+        }
+        """
+        return dict(query=mutation, variables={"alerts": self.alerts})
+
+
 def get_alerts(dict):
     alerts_api_url = get_forta_api_url()
     headers = get_forta_api_headers()
diff --git a/python-sdk/src/forta_agent/finding.py b/python-sdk/src/forta_agent/finding.py
index 77b3c2b..d81a1a0 100644
--- a/python-sdk/src/forta_agent/finding.py
+++ b/python-sdk/src/forta_agent/finding.py
@@ -1,4 +1,5 @@
 import json
+from datetime import datetime
 from .label import Label
 from enum import IntEnum
 from .utils import assert_enum_value_in_dict, assert_non_empty_string_in_dict
@@ -51,6 +52,7 @@ def __init__(self, dict):
             l, Label) else Label(l), dict.get('labels', [])))
         self.unique_key = dict.get('unique_key')
         self.source = dict.get('source')
+        self.timestamp = dict.get('timestamp', datetime.now())
 
     def toJson(self):
         d = dict(self.__dict__, **{
@@ -58,4 +60,4 @@ def toJson(self):
             'labels': list(map(lambda l: l.toDict(), self.labels)),
             'uniqueKey': self.unique_key
         })
-        return json.dumps({k: v for k, v in d.items() if v or k == 'type' or k == 'severity'})
+        return json.dumps({k: v for k, v in d.items() if v or k == 'type' or k == 'severity'}, default=str)
diff --git a/python-sdk/src/forta_agent/utils.py b/python-sdk/src/forta_agent/utils.py
index 4a67e2c..c564f76 100644
--- a/python-sdk/src/forta_agent/utils.py
+++ b/python-sdk/src/forta_agent/utils.py
@@ -98,6 +98,8 @@ def get_forta_api_headers():
     config = get_forta_config()
     if "fortaApiKey" in config:
         headers["Authorization"] = f'Bearer {config.get("fortaApiKey")}'
+    elif 'FORTA_API_KEY' in os.environ:
+        headers["Authorization"] = f'Bearer {os.environ["FORTA_API_KEY"]}'
 
     return headers