Merge pull request #261 from forta-network/add-labels-to-alerts

add labels to alerts
forta-network · Feb 14, 2023 · b8bf239 · b8bf239
2 parents dd0b72d + c28a784
commit b8bf239
Show file tree

Hide file tree

Showing 9 changed files with 44 additions and 4 deletions.
diff --git a/cli/commands/run/server/agent.controller.js b/cli/commands/run/server/agent.controller.js
@@ -3,6 +3,7 @@ const {
  TransactionEvent,
  isPrivateFindings,
  AlertEvent,
+ Label
 } = require("../../../../sdk");
 const {
  assertExists,
@@ -188,6 +189,7 @@ module.exports = class AgentController {
  createAlertEventFromGrpcRequest(request) {
  const { alert } = request.event;
 
+ alert.labels = alert.labels.map((l) => Label.fromObject(l));
  return new AlertEvent(alert);
  }
 

diff --git a/cli/commands/run/server/agent.controller.spec.ts b/cli/commands/run/server/agent.controller.spec.ts
@@ -53,7 +53,8 @@ describe("AgentController", () => {
  request: {
  event: {
  alert: {
- hash: "0x123"
+ hash: "0x123",
+ labels: []
  }
  }
  }

diff --git a/python-sdk/src/forta_agent/alert.py b/python-sdk/src/forta_agent/alert.py
@@ -1,5 +1,6 @@
 class Alert:
  def __init__(self, dict):
+ from .label import Label
  self.addresses = dict.get('addresses')
  self.alert_id = dict.get('alertId')
  self.contracts = list(map(lambda t: Contract(t), dict.get(
@@ -20,6 +21,8 @@ def __init__(self, dict):
  self.alert_document_type = dict.get('alertDocumentType')
  self.related_alerts = dict.get('relatedAlerts')
  self.chain_id = dict.get('chainId')
+ self.labels = list(map(lambda t: Label(t), dict.get(
+ 'labels', []))) if dict.get('labels') is not None else []
 
 
 class Source:

diff --git a/python-sdk/src/forta_agent/forta_graphql.py b/python-sdk/src/forta_agent/forta_graphql.py
@@ -84,6 +84,13 @@ def get_query(self):
  findingType
  relatedAlerts
  chainId
+ labels {
+ label
+ confidence
+ entity
+ entityType
+ remove
+ }
  }
  pageInfo {
  hasNextPage

diff --git a/python-sdk/src/forta_agent/label.py b/python-sdk/src/forta_agent/label.py
@@ -12,12 +12,13 @@ class EntityType(IntEnum):
 
 class Label:
  def __init__(self, dict):
- assert_enum_value_in_dict(dict, 'entity_type', EntityType)
+ entityTypeVal = dict.get('entity_type', dict['entityType'])
+ self.entity_type = EntityType[entityTypeVal.title()] if type(
+ entityTypeVal) == str else EntityType(entityTypeVal)
+ assert_enum_value_in_dict(self.__dict__, 'entity_type', EntityType)
  assert_non_empty_string_in_dict(dict, 'entity')
- assert_non_empty_string_in_dict(dict, 'label')
  self.entity = dict['entity']
  self.confidence = dict['confidence']
- self.entity_type = dict['entity_type']
  self.label = dict['label']
  self.remove = dict.get('remove', False)
 

diff --git a/sdk/alert.ts b/sdk/alert.ts
@@ -1,3 +1,5 @@
+import { Label } from "./label"
+
 export interface Alert {
  addresses?: string[],
  alertId?: string,
@@ -17,6 +19,7 @@ export interface Alert {
  alertDocumentType?: string,
  relatedAlerts?: string[],
  chainId?: number,
+ labels?: Label[],
  source?: {
  transactionHash?: string,
  block?: {

diff --git a/sdk/graphql/forta.ts b/sdk/graphql/forta.ts
@@ -154,6 +154,13 @@ export const getQueryFromAlertOptions = (options: AlertQueryOptions) => {
  findingType
  relatedAlerts
  chainId
+ labels {
+ label
+ confidence
+ entity
+ entityType
+ remove
+ }
  }
  pageInfo {
  hasNextPage

diff --git a/sdk/label.ts b/sdk/label.ts
@@ -6,6 +6,14 @@ export enum EntityType {
  Url,
 }
 
+export const ENTITY_TYPE_STRING_TO_ENUM = {
+ UNKNOWN: EntityType.Unknown,
+ ADDRESS: EntityType.Address,
+ TRANSACTION: EntityType.Transaction,
+ BLOCK: EntityType.Block,
+ URL: EntityType.Url,
+};
+
 type LabelInput = {
  entityType: EntityType;
  entity: string;
@@ -30,6 +38,9 @@ export class Label {
  confidence,
  remove = false,
  }: LabelInput) {
+ if (typeof entityType == "string") {
+ entityType = ENTITY_TYPE_STRING_TO_ENUM[entityType];
+ }
  return new Label(entityType, entity, label, confidence, remove);
  }
 }
diff --git a/sdk/utils.ts b/sdk/utils.ts
@@ -10,6 +10,7 @@ import {
  BlockEvent,
  EventType,
  FortaConfig,
+ Label,
  Network,
  Trace,
  TransactionEvent
@@ -173,5 +174,9 @@ export const getAlerts = async (query: AlertQueryOptions): Promise<AlertsRespons
 
  if(response.data && response.data.errors) throw Error(response.data.errors)
 
+ const alerts = response.data.data.alerts.alerts
+ for (const alert of alerts) {
+ alert.labels = alert.labels!.map((l) => Label.fromObject(l));
+ }
  return response.data.data.alerts
 }