psengine 2.0.4__py3-none-any.whl

psengine/constants.py

@@ -0,0 +1,63 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+import sys
+
+#####################
+# Multithreading Defaults
+#####################
+DEFAULT_LIMIT = 10
+DEFAULT_MAX_WORKERS = 10
+
+#####################
+# Recorded Future API
+#####################
+RF_TOKEN_ENV_VAR = 'RF_TOKEN'
+RF_TOKEN_VALIDATION_REGEX = r'^[a-f0-9]{32}$'
+
+#####################
+# Recorded Future Portal
+#####################
+RF_PORTAL_BASE_URL = 'https://app.recordedfuture.com'
+
+# 0: entity type, 1: entity name
+INDICATOR_INTEL_CARD_URL = RF_PORTAL_BASE_URL + '/live/sc/entity/{}:{}'
+
+#####################
+# HTTP Client Defaults
+#####################
+# In seconds
+REQUEST_TIMEOUT = 120
+POOL_MAX_SIZE = 120
+
+# Request RETRY configuration
+RETRY_TOTAL = 5
+BACKOFF_FACTOR = 1
+STATUS_FORCELIST = [502, 503, 504]
+
+SSL_VERIFY = True
+
+#####################
+# Misc
+#####################
+# Output
+ROOT_DIR = sys.path[0]
+
+# String for timestamp conversion
+TIMESTAMP_STR = '%Y-%m-%d %H:%M:%S'
+
+# Markdown truncation string
+TRUNCATE_COMMENT = (
+    '\n\nThis {type_} has been truncated due to a character limit imposed by this tool. '
+    'View the full {type_} at {url}'
+)

psengine/detection/__init__.py

@@ -0,0 +1,17 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from .detection_mgr import DetectionMgr
+from .detection_rule import DetectionRule, DetectionRuleSearchOut
+from .errors import DetectionRuleFetchError, DetectionRuleSearchError
+from .helpers import save_rule

psengine/detection/detection_mgr.py

@@ -0,0 +1,135 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+import logging
+from typing import Optional, Union
+
+from pydantic import validate_call
+
+from ..constants import DEFAULT_LIMIT
+from ..endpoints import EP_DETECTION_RULES
+from ..helpers import debug_call
+from ..helpers.helpers import connection_exceptions
+from ..rf_client import RFClient
+from .detection_rule import DetectionRule, DetectionRuleSearchOut
+from .errors import DetectionRuleFetchError, DetectionRuleSearchError
+
+SEARCH_LIMIT = 100
+
+
+class DetectionMgr:
+    """Class to manage DetectionRules and interaction with the Detection API."""
+
+    def __init__(self, rf_token: str = None):
+        """Initializes the DetectionMgr object.
+
+        Args:
+            rf_token (str, optional): Recorded Future API token. Defaults to None
+        """
+        self.log = logging.getLogger(__name__)
+        self.rf_client = RFClient(api_token=rf_token) if rf_token else RFClient()
+
+    @debug_call
+    @validate_call
+    @connection_exceptions(ignore_status_code=[], exception_to_raise=DetectionRuleSearchError)
+    def search(
+        self,
+        detection_rule: Union[list[str], str, None] = None,
+        entities: Optional[list[str]] = None,
+        created_before: Optional[str] = None,
+        created_after: Optional[str] = None,
+        updated_before: Optional[str] = None,
+        updated_after: Optional[str] = None,
+        doc_id: Optional[str] = None,
+        title: Optional[str] = None,
+        tagged_entities: Optional[bool] = None,
+        max_results: Optional[int] = DEFAULT_LIMIT,
+    ) -> list[DetectionRule]:
+        """Search for detection rules based on various filter criteria.
+
+        Endpoint:
+            ``detection-rule/search``
+
+        Args:
+            detection_rule: Types of detection rules to search for. Defaults to None.
+            entities: List of entities to filter the search. Defaults to None.
+            created_before: Filter for rules created before this date. Defaults to None.
+            created_after: Filter for rules created after this date. Defaults to None.
+            updated_before: Filter for rules updated before this date. Defaults to None.
+            updated_after: Filter for rules updated after this date. Defaults to None.
+            doc_id: Filter by document ID. Defaults to None.
+            title: Filter by title. Defaults to None.
+            tagged_entities: Whether to filter by tagged entities. Defaults to None.
+            max_results: Limit the total number of results returned. Defaults to 10.
+
+        Raises:
+            ValidationError if any supplied parameter is of incorrect type.
+            DetectionRuleSearchError: if connection error occurs.
+
+        Returns:
+            List[DetectionRule]: A list of detection rules matching the search criteria.
+        """
+        detection_rule = [detection_rule] if isinstance(detection_rule, str) else detection_rule
+        filters = {
+            'types': detection_rule,
+            'entities': entities,
+            'created': {'before': created_before, 'after': created_after},
+            'updated': {'before': updated_before, 'after': updated_after},
+            'doc_id': doc_id,
+            'title': title,
+        }
+        data = {
+            'filter': filters,
+            'tagged_entities': tagged_entities,
+            'limit': SEARCH_LIMIT,
+        }
+
+        data = DetectionRuleSearchOut.model_validate(data)
+        results = self.rf_client.request_paged(
+            'post',
+            EP_DETECTION_RULES,
+            data=data.json(),
+            results_path='result',
+            offset_key='offset',
+            max_results=max_results,
+        )
+
+        results = results if isinstance(results, list) else results.json().get('result', [])
+        return [DetectionRule.model_validate(data) for data in results]
+
+    @debug_call
+    @validate_call
+    def fetch(self, doc_id: str) -> Optional[DetectionRule]:
+        """Fetch of a detection rule based on rule id.
+
+        Endpoint:
+            ``detection-rule/search``
+
+        Args:
+            doc_id (str): Detection rule id to lookup.
+
+        Raises:
+            ValidationError if any supplied parameter is of incorrect type.
+            DetectionRuleLookupError: If no rule is found for the given id.
+
+        Returns:
+            Optional[DetectionRule]: The detection rule found for the given id.
+        """
+        try:
+            result = self.search(doc_id=doc_id)
+        except DetectionRuleSearchError as e:
+            raise DetectionRuleFetchError(f'Error in fething of {doc_id}') from e
+
+        if result:
+            return result[0]
+        self.log.info(f'No rule found for id {doc_id}')

psengine/detection/detection_rule.py

@@ -0,0 +1,85 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from datetime import datetime
+from functools import total_ordering
+from typing import Optional
+
+from pydantic import Field
+
+from ..common_models import RFBaseModel
+from ..constants import TIMESTAMP_STR
+from .models import DetectionRuleType, RuleContext, SearchFilter
+
+
+@total_ordering
+class DetectionRule(RFBaseModel):
+    """Detection rule model to validate output of the ``/search`` endpoint.
+
+    Methods:
+        __hash__:
+            Returns a hash value based on ``id_`` and updated timestamp.
+
+        __eq__:
+            Checks equality between two DetectionRule instances based on ``id_`` and updated time.
+
+        __gt__:
+            Defines a greater-than comparison between two DetectionRule instances based on
+            updated timestamp and ``id_``.
+
+        __str__:
+            Returns a string representation of the DetectionRule instance with:
+            ``id_``, created timestamp, updated timestamp, and title.
+
+            .. code-block:: python
+
+                >>> print(detection_rule)
+                ID: rule123, Created: 2024-05-21 10:42:30AM, Updated: 2024-05-21 10:42:30AM,
+                Title: Example Rule
+
+    Total Ordering:
+        The ordering of DetectionRule instances is determined primarily by the updated timestamp.
+        If two instances have the same updated timestamp, ``id_`` is used as a secondary criterion.
+    """
+
+    id_: str = Field(alias='id')
+    type_: DetectionRuleType = Field(alias='type')
+    title: str
+    description: str
+    created: datetime
+    updated: datetime
+    rules: list[RuleContext]
+
+    def __hash__(self):
+        return hash((self.id_, self.updated))
+
+    def __eq__(self, other: 'DetectionRule'):
+        return (self.id_, self.updated) == (other.id_, other.updated)
+
+    def __gt__(self, other: 'DetectionRule'):
+        return (self.updated, self.id_) > (other.updated, other.id_)
+
+    def __str__(self):
+        return (
+            f'ID: {self.id_}, Created: {self.created.strftime(TIMESTAMP_STR)}, '
+            f'Updated: {self.updated.strftime(TIMESTAMP_STR)}, Title: {self.title}'
+        )
+
+
+class DetectionRuleSearchOut(RFBaseModel):
+    """Model to validate ``/search`` endpoint payload sent."""
+
+    filter_: Optional[SearchFilter] = Field(alias='filter', default={})
+    tagged_entities: Optional[bool] = False
+    limit: Optional[int] = None
+    offset: Optional[str] = None

psengine/detection/errors.py

@@ -0,0 +1,26 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from ..errors import RecordedFutureError
+
+
+class DetectionRulesError(RecordedFutureError):
+    """Error raised when there was an error with the detection rules API."""
+
+
+class DetectionRuleSearchError(DetectionRulesError):
+    """Error raised when there was an error searching for detection rules."""
+
+
+class DetectionRuleFetchError(DetectionRulesError):
+    """Error raised when there was an error fetching of detection rule id."""

psengine/detection/helpers.py

@@ -0,0 +1,56 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+import logging
+from pathlib import Path
+from typing import Union
+
+from pydantic import validate_call
+
+from ..errors import WriteFileError
+from ..helpers import OSHelpers, debug_call
+from .detection_rule import DetectionRule
+
+LOG = logging.getLogger('psengine.detection.helpers')
+
+
+@debug_call
+@validate_call
+def save_rule(rule: DetectionRule, output_directory: Union[str, Path] = None):
+    """Write detection rule content to file. If more than one detection rule is attached to rule,
+    all will be saved.
+
+    Args:
+        rule (DetectionRule): single detection to write.
+        output_directory (Union[str, Path]): a path to write to. If not provided, it will be the
+        current working directory.
+
+    Raises:
+       WriteFileError: if the path provided is not a directory or it cannot be created.
+       WriteFileError: if the write operations fail.
+
+    """
+    if not rule.rules:
+        LOG.info(f'No rules to write for {rule.id_}')
+        return
+
+    output_directory = Path(output_directory).absolute() if output_directory else Path().cwd()
+    OSHelpers.mkdir(output_directory)
+
+    for data in rule.rules:
+        try:
+            full_path = output_directory / data.file_name
+            full_path.write_text(data.content)
+            LOG.info(f'Wrote: {full_path}')
+        except (FileNotFoundError, IsADirectoryError, PermissionError, OSError) as err:  # noqa: PERF203
+            raise WriteFileError(f"Could not write file '{data.file_name}': {err}") from err

psengine/detection/models.py

@@ -0,0 +1,47 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import Field
+
+from ..common_models import DetectionRuleType, RFBaseModel
+
+
+class Entity(RFBaseModel):
+    id_: str = Field(alias='id', default=None)
+    name: Optional[str] = None
+    type_: str = Field(alias='type', default=None)
+
+    display_name: Optional[str] = None
+
+
+class RuleContext(RFBaseModel):
+    entities: list[Entity]
+    content: str
+    file_name: Optional[str] = None
+
+
+class TimeRange(RFBaseModel):
+    after: Optional[datetime] = None
+    before: Optional[datetime] = None
+
+
+class SearchFilter(RFBaseModel):
+    types: Optional[list[DetectionRuleType]] = None
+    entities: Optional[list[str]] = None
+    created: Optional[TimeRange] = None
+    updated: Optional[TimeRange] = None
+    doc_id: Optional[str] = None
+    title: Optional[str] = None

psengine/endpoints.py

@@ -0,0 +1,98 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+from os import environ
+
+###############################################################################
+# API Version
+###############################################################################
+API_VERSION = 'v2'
+
+###############################################################################
+# API Base URLs
+###############################################################################
+BASE_URL = 'https://api.recordedfuture.com'
+CONNECT_API_BASE_URL = BASE_URL + '/' + API_VERSION
+
+BASE_URL = environ.get('RF_BASE_URL') if environ.get('RF_BASE_URL') else BASE_URL
+CONNECT_API_BASE_URL = (
+    environ.get('RF_BASE_URL') if environ.get('RF_BASE_URL') else CONNECT_API_BASE_URL
+)
+
+###############################################################################
+# Classic Alerts Endpoints V3
+###############################################################################
+EP_CLASSIC_ALERTS_V3 = BASE_URL + '/v3'
+EP_CLASSIC_ALERTS_RULES = CONNECT_API_BASE_URL + '/alert/rule'
+EP_CLASSIC_ALERTS_UPDATE = CONNECT_API_BASE_URL + '/alert/update'
+EP_CLASSIC_ALERTS_SEARCH = EP_CLASSIC_ALERTS_V3 + '/alerts/'
+EP_CLASSIC_ALERTS_HITS = EP_CLASSIC_ALERTS_V3 + '/alerts/hits'
+EP_CLASSIC_ALERTS_IMAGE = EP_CLASSIC_ALERTS_V3 + '/alerts/image'
+EP_CLASSIC_ALERTS_ID = EP_CLASSIC_ALERTS_V3 + '/alerts/{}'
+
+###############################################################################
+# Fusion Endpoints
+###############################################################################
+EP_RISKLIST = f'{BASE_URL}/{API_VERSION}/' + '{}/risklist'
+EP_FUSION_FILES = CONNECT_API_BASE_URL + '/fusion/files'
+
+###############################################################################
+# Playbook Alert Endpoints
+###############################################################################
+EP_PLAYBOOK_ALERT = BASE_URL + '/playbook-alert'
+EP_PLAYBOOK_ALERT_SEARCH = EP_PLAYBOOK_ALERT + '/search'
+EP_PLAYBOOK_ALERT_COMMON = EP_PLAYBOOK_ALERT + '/common'
+EP_PLAYBOOK_ALERT_DOMAIN_ABUSE = EP_PLAYBOOK_ALERT + '/domain_abuse'
+EP_PLAYBOOK_ALERT_CYBER_VULNERABILITY = EP_PLAYBOOK_ALERT + '/vulnerability'
+EP_PLAYBOOK_ALERT_CODE_REPO_LEAKAGE = EP_PLAYBOOK_ALERT + '/code_repo_leakage'
+EP_PLAYBOOK_ALERT_THIRD_PARTY_RISK = EP_PLAYBOOK_ALERT + '/third_party_risk'
+EP_PLAYBOOK_ALERT_IDENTITY_NOVEL_EXPOSURES = EP_PLAYBOOK_ALERT + '/identity_novel_exposures'
+
+###############################################################################
+# Entity Match Endpoint
+###############################################################################
+EP_ENTITY_MATCH = BASE_URL + '/entity-match/match'
+EP_ENTITY_LOOKUP = BASE_URL + '/entity-match/entity/{}'
+
+###############################################################################
+# List API Endpoints
+###############################################################################
+EP_LIST = BASE_URL + '/list'
+EP_CREATE_LIST = EP_LIST + '/create'
+EP_SEARCH_LIST = EP_LIST + '/search'
+
+###############################################################################
+# SOAR Endpoints
+###############################################################################
+EP_SOAR_ENRICHMENT = CONNECT_API_BASE_URL + '/soar/enrichment'
+
+###############################################################################
+# Detection Rules API Endpoints
+###############################################################################
+EP_DETECTION_RULES = BASE_URL + '/detection-rule/search'
+
+###############################################################################
+# Collective Insights API Endpoints
+###############################################################################
+EP_COLLECTIVE_INSIGHTS = BASE_URL + '/collective-insights'
+EP_COLLECTIVE_INSIGHTS_DETECTIONS = EP_COLLECTIVE_INSIGHTS + '/detections'
+
+###############################################################################
+# Analyst Notes API Endpoints
+###############################################################################
+EP_ANALYST_NOTE = BASE_URL + '/analyst-note/'
+EP_ANALYST_NOTE_SEARCH = EP_ANALYST_NOTE + 'search'
+EP_ANALYST_NOTE_LOOKUP = EP_ANALYST_NOTE + 'lookup/{}'
+EP_ANALYST_NOTE_PREVIEW = EP_ANALYST_NOTE + 'preview'
+EP_ANALYST_NOTE_PUBLISH = EP_ANALYST_NOTE + 'publish'
+EP_ANALYST_NOTE_DELETE = EP_ANALYST_NOTE + 'delete/{}'
+EP_ANALYST_NOTE_ATTACHMENT = EP_ANALYST_NOTE + 'attachment/{}'

psengine/enrich/__init__.py

@@ -0,0 +1,28 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+
+from .errors import EnrichmentLookupError, EnrichmentSoarError
+from .lookup import (
+    EnrichedCompany,
+    EnrichedDomain,
+    EnrichedHash,
+    EnrichedIP,
+    EnrichedMalware,
+    EnrichedURL,
+    EnrichedVulnerability,
+    EnrichmentData,
+)
+from .lookup_mgr import LookupMgr
+from .soar import SOAREnrichedEntity, SOAREnrichIn, SOAREnrichOut
+from .soar_mgr import SoarMgr

psengine/enrich/constants.py

@@ -0,0 +1,73 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from typing import Literal
+
+from ..enrich import (
+    EnrichedCompany,
+    EnrichedDomain,
+    EnrichedHash,
+    EnrichedIP,
+    EnrichedMalware,
+    EnrichedURL,
+    EnrichedVulnerability,
+)
+
+SOAR_POST_ROWS = 2000
+
+ALLOWED_ENTITIES = Literal[
+    'company',
+    'company_by_domain',
+    'company/by_domain',
+    'domain',
+    'hash',
+    'ip',
+    'malware',
+    'url',
+    'vulnerability',
+    'Company',
+    'Organization',
+    'InternetDomainName',
+    'Hash',
+    'IpAddress',
+    'Malware',
+    'URL',
+    'CyberVulnerability',
+]
+
+ENTITY_FIELDS = ['entity', 'risk', 'timestamps']
+MALWARE_FIELDS = ['entity', 'timestamps']
+TYPE_MAPPING = {
+    'company/by_domain': 'company_by_domain',
+    'Organization': 'company',
+    'Company': 'company',
+    'IpAddress': 'ip',
+    'InternetDomainName': 'domain',
+    'CyberVulnerability': 'vulnerability',
+    'URL': 'url',
+    'Malware': 'malware',
+    'Hash': 'hash',
+}
+
+
+IOC_TO_MODEL = {
+    'ip': EnrichedIP,
+    'domain': EnrichedDomain,
+    'hash': EnrichedHash,
+    'url': EnrichedURL,
+    'malware': EnrichedMalware,
+    'vulnerability': EnrichedVulnerability,
+    'company': EnrichedCompany,
+    'company/by_domain': EnrichedCompany,
+}
+MESSAGE_404 = '404 received. Nothing known on this entity'

psengine/enrich/errors.py

@@ -0,0 +1,26 @@
+##################################### TERMS OF USE ###########################################
+# The following code is provided for demonstration purpose only, and should not be used      #
+# without independent verification. Recorded Future makes no representations or warranties,  #
+# express, implied, statutory, or otherwise, regarding any aspect of this code or of the     #
+# information it may retrieve, and provides it both strictly “as-is” and without assuming    #
+# responsibility for any information it may retrieve. Recorded Future shall not be liable    #
+# for, and you assume all risk of using, the foregoing. By using this code, Customer         #
+# represents that it is solely responsible for having all necessary licenses, permissions,   #
+# rights, and/or consents to connect to third party APIs, and that it is solely responsible  #
+# for having all necessary licenses, permissions, rights, and/or consents to any data        #
+# accessed from any third party API.                                                         #
+##############################################################################################
+
+from ..errors import RecordedFutureError
+
+
+class EnrichmentError(RecordedFutureError):
+    """Error raised when there was an error enriching an entity via Lookup or SOAR."""
+
+
+class EnrichmentLookupError(EnrichmentError):
+    """Error raised when there was an error looking up for an entity name or entity id."""
+
+
+class EnrichmentSoarError(EnrichmentError):
+    """Error raised when there was an error performing a SOAR enrichment."""