Skip to content

Commit

Permalink
python external bot support
Browse files Browse the repository at this point in the history
  • Loading branch information
haseebrabbani committed Oct 12, 2023
1 parent f9d2d48 commit 3f52514
Show file tree
Hide file tree
Showing 4 changed files with 76 additions and 2 deletions.
2 changes: 1 addition & 1 deletion python-sdk/src/forta_agent/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
70 changes: 70 additions & 0 deletions python-sdk/src/forta_agent/alerts_api.py
Original file line number Diff line number Diff line change
@@ -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()
Expand Down
4 changes: 3 additions & 1 deletion python-sdk/src/forta_agent/finding.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -51,11 +52,12 @@ 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__, **{
'alertId': self.alert_id,
'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)
2 changes: 2 additions & 0 deletions python-sdk/src/forta_agent/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit 3f52514

Please sign in to comment.