diff --git a/core/schemas/indicator.py b/core/schemas/indicator.py index 9ac237323..1b7aed69a 100644 --- a/core/schemas/indicator.py +++ b/core/schemas/indicator.py @@ -118,22 +118,13 @@ def match(self, value: str) -> IndicatorMatch | None: return None -class QueryType(str, Enum): - opensearch = "opensearch" - osquery = "osquery" - sql = "sql" - splunk = "splunk" - censys = "censys" - shodan = "shodan" - - class Query(Indicator): """Represents a query that can be sent to another system.""" _type_filter: ClassVar[str] = IndicatorType.query type: Literal["query"] = IndicatorType.query - query_type: QueryType + query_type: str target_systems: list[str] = [] def match(self, value: str) -> IndicatorMatch | None: diff --git a/plugins/analytics/public/censys.py b/plugins/analytics/public/censys.py index 6d7d3675c..5c51a03a3 100644 --- a/plugins/analytics/public/censys.py +++ b/plugins/analytics/public/censys.py @@ -31,9 +31,7 @@ def run(self): api_secret=api_secret, ) - censys_queries, _ = indicator.Query.filter( - {"query_type": indicator.QueryType.censys} - ) + censys_queries, _ = indicator.Query.filter({"query_type": "censys"}) for query in censys_queries: ip_addresses = query_censys(hosts_api, query.pattern) diff --git a/plugins/analytics/public/shodan.py b/plugins/analytics/public/shodan.py index d8fa2d6f1..e74b9f533 100644 --- a/plugins/analytics/public/shodan.py +++ b/plugins/analytics/public/shodan.py @@ -28,9 +28,7 @@ def run(self): shodan_api = Shodan(api_key) - shodan_queries, _ = indicator.Query.filter( - {"query_type": indicator.QueryType.shodan} - ) + shodan_queries, _ = indicator.Query.filter({"query_type": "shodan"}) for query in shodan_queries: ip_addresses = query_shodan(shodan_api, query.pattern, result_limit) diff --git a/plugins/feeds/public/dfiq.py b/plugins/feeds/public/dfiq.py index 00ea96a77..90eadfbdb 100644 --- a/plugins/feeds/public/dfiq.py +++ b/plugins/feeds/public/dfiq.py @@ -52,7 +52,7 @@ def _process_approach(yaml_string: str) -> None: name=step.description, pattern=step.value, relevant_tags=approach.dfiq_tags or [], - query_type=indicator.QueryType.opensearch, + query_type="opensearch", location=processor.name, diamond=indicator.DiamondModel.victim, ).save() diff --git a/tests/analytics_test.py b/tests/analytics_test.py index e816d2eb2..a5376b68a 100644 --- a/tests/analytics_test.py +++ b/tests/analytics_test.py @@ -78,7 +78,7 @@ def test_censys_query(self, mock_censys_hosts): location="censys", diamond=DiamondModel.infrastructure, relevant_tags=["censys_query_tag"], - query_type=indicator.QueryType.censys, + query_type="censys", ).save() mock_search_result = [ @@ -144,7 +144,7 @@ def test_shodan_query_with_various_limits(self, limit, expected_count, mock_shod location="shodan", diamond=DiamondModel.infrastructure, relevant_tags=["shodan_query_tag"], - query_type=indicator.QueryType.shodan, + query_type="shodan", ).save() def mock_search_cursor(query): @@ -190,7 +190,7 @@ def test_shodan_observables_and_neighbors(self, mock_shodan): location="shodan", diamond=DiamondModel.infrastructure, relevant_tags=["shodan_query_tag"], - query_type=indicator.QueryType.shodan, + query_type="shodan", ).save() def mock_search_cursor(query): diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py index d08718a83..950298ee0 100644 --- a/tests/schemas/fixture.py +++ b/tests/schemas/fixture.py @@ -2,7 +2,7 @@ from core import database_arango from core.schemas.entity import Investigation, Malware, ThreatActor -from core.schemas.indicator import DiamondModel, Query, QueryType, Regex +from core.schemas.indicator import DiamondModel, Query, Regex from core.schemas.observables import ( bic, generic_observable, @@ -79,7 +79,7 @@ def test_something(self): location="syslogs", diamond=DiamondModel.capability, pattern='(reporter:"sshd" AND Accepted)', - query_type=QueryType.opensearch, + query_type="opensearch", target_systems=["timesketch", "plaso"], relevant_tags=["ssh", "login"], ).save()