Skip to content

Commit f79f1aa

Browse files
nor3thnor3th
authored
[examples/tests/api] Added enrichment example, OpenCTIApiWork improvements and further test cases (OpenCTI-Platform#202)
* [examples] Added ask_enrichment of Observable example * [tests] Added OpenCTIApiWork as fixture * [pycti] Return work_id after asking for SDO enrichment * [pycti/tests] Migrated work related functions from test utils to OpenCTIApiWork * [tests] Folder restructuring * [api] Work by id query implemented * [stix2] Added mapping for integer value SimpleObservables * [stix2] Convert SimpleObservable values to int if the key is in the INT_VALUES list * [tests] Created simple test case for STIX object import * [tests] Small refactoring * [tests] Added fixtures to STIX tests * [tests] Changed ownclass, baseclass and stixclass to snake case and format adaptions
1 parent 976f39e commit f79f1aa

17 files changed

+512
-332
lines changed

‎examples/ask_enrichment_of_observable.py

+47
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# coding: utf-8
2+
from pycti import OpenCTIApiClient, OpenCTIApiConnector, OpenCTIApiWork
3+
4+
# Variables
5+
api_url = "https://demo.opencti.io"
6+
api_token = "YOUR_TOKEN"
7+
# Define name of INTERNAL_ENRICHMENT Connector which can enrich IPv4 addresses
8+
connector_name = "AbuseIPDB"
9+
10+
# OpenCTI initialization
11+
opencti_api_client = OpenCTIApiClient(api_url, api_token)
12+
opencti_api_connector = OpenCTIApiConnector(opencti_api_client)
13+
opencti_api_work = OpenCTIApiWork(opencti_api_client)
14+
15+
# Create the observable
16+
observable = opencti_api_client.stix_cyber_observable.create(
17+
**{
18+
"simple_observable_key": "IPv4-Addr.value",
19+
"simple_observable_value": "8.8.4.4",
20+
}
21+
)
22+
23+
# Get connector id for defined connector name
24+
connector_list = opencti_api_connector.list()
25+
connector_names = []
26+
connector_id = ""
27+
for connector in connector_list:
28+
connector_names.append(connector["name"])
29+
if connector["name"] == connector_name:
30+
connector_id = connector["id"]
31+
32+
if connector_id == "":
33+
print(f"Connector with name '{connector_name}' could not be found")
34+
print(f"Running connectors: {connector_names}")
35+
exit(0)
36+
37+
print("Asking for enrichment... (this might take a bit to finish)")
38+
# Ask for enrichment
39+
work_id = opencti_api_client.stix_cyber_observable.ask_for_enrichment(
40+
id=observable["id"], connector_id=connector_id
41+
)
42+
# Wait for connector to finish
43+
opencti_api_work.wait_for_work_to_finish(work_id)
44+
45+
# Read the observable
46+
obs = opencti_api_client.stix_cyber_observable.read(id=observable["id"])
47+
print(obs)

‎pycti/api/opencti_api_work.py

+137
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import logging
2+
import time
3+
from typing import Dict, List
24

35

46
class OpenCTIApiWork:
@@ -72,3 +74,138 @@ def initiate_work(self, connector_id: str, friendly_name: str) -> str:
7274
query, {"connectorId": connector_id, "friendlyName": friendly_name}
7375
)
7476
return work["data"]["workAdd"]["id"]
77+
78+
def delete_work(self, work_id: str):
79+
query = """
80+
mutation ConnectorWorksMutation($workId: ID!) {
81+
workEdit(id: $workId) {
82+
delete
83+
}
84+
}"""
85+
work = self.api.query(
86+
query,
87+
{"workId": work_id},
88+
)
89+
return work["data"]
90+
91+
def wait_for_work_to_finish(self, work_id: str):
92+
status = ""
93+
cnt = 0
94+
while status != "complete":
95+
state = self.get_work(work_id=work_id)
96+
if len(state) > 0:
97+
status = state["status"]
98+
99+
if state["errors"]:
100+
self.api.log(
101+
"error", f"Unexpected connector error {state['errors']}"
102+
)
103+
return ""
104+
105+
time.sleep(1)
106+
cnt += 1
107+
108+
def get_work(self, work_id: str) -> Dict:
109+
query = """
110+
query WorkQuery($id: ID!) {
111+
work(id: $id) {
112+
id
113+
name
114+
user {
115+
name
116+
}
117+
timestamp
118+
status
119+
event_source_id
120+
received_time
121+
processed_time
122+
completed_time
123+
tracking {
124+
import_expected_number
125+
import_processed_number
126+
}
127+
messages {
128+
timestamp
129+
message
130+
sequence
131+
source
132+
}
133+
errors {
134+
timestamp
135+
message
136+
sequence
137+
source
138+
}
139+
}
140+
}
141+
"""
142+
result = self.api.query(
143+
query,
144+
{"id": work_id},
145+
)
146+
return result["data"]["work"]
147+
148+
def get_connector_works(self, connector_id: str) -> List[Dict]:
149+
query = """
150+
query ConnectorWorksQuery(
151+
$count: Int
152+
$orderBy: WorksOrdering
153+
$orderMode: OrderingMode
154+
$filters: [WorksFiltering]
155+
) {
156+
works(
157+
first: $count
158+
orderBy: $orderBy
159+
orderMode: $orderMode
160+
filters: $filters
161+
) {
162+
edges {
163+
node {
164+
id
165+
name
166+
user {
167+
name
168+
}
169+
timestamp
170+
status
171+
event_source_id
172+
received_time
173+
processed_time
174+
completed_time
175+
tracking {
176+
import_expected_number
177+
import_processed_number
178+
}
179+
messages {
180+
timestamp
181+
message
182+
sequence
183+
source
184+
}
185+
errors {
186+
timestamp
187+
message
188+
sequence
189+
source
190+
}
191+
}
192+
}
193+
}
194+
}
195+
"""
196+
result = self.api.query(
197+
query,
198+
{
199+
"count": 50,
200+
"filters": [
201+
{"key": "connector_id", "values": [connector_id]},
202+
],
203+
},
204+
)
205+
result = result["data"]["works"]["edges"]
206+
return_value = []
207+
for node in result:
208+
node = node["node"]
209+
return_value.append(node)
210+
211+
return sorted(return_value, key=lambda i: i["timestamp"])

‎pycti/entities/opencti_stix_cyber_observable.py

+5-4
Original file line numberDiff line numberDiff line change
@@ -1728,13 +1728,13 @@ def push_list_export(self, file_name, data, list_filters=""):
17281728
},
17291729
)
17301730

1731-
def ask_for_enrichment(self, **kwargs):
1731+
def ask_for_enrichment(self, **kwargs) -> str:
17321732
id = kwargs.get("id", None)
17331733
connector_id = kwargs.get("connector_id", None)
17341734

17351735
if id is None or connector_id is None:
17361736
self.opencti.log("error", "Missing parameters: id and connector_id")
1737-
return False
1737+
return ""
17381738

17391739
query = """
17401740
mutation StixCoreObjectEnrichmentLinesMutation($id: ID!, $connectorId: ID!) {
@@ -1746,11 +1746,12 @@ def ask_for_enrichment(self, **kwargs):
17461746
}
17471747
"""
17481748

1749-
self.opencti.query(
1749+
result = self.opencti.query(
17501750
query,
17511751
{
17521752
"id": id,
17531753
"connectorId": connector_id,
17541754
},
17551755
)
1756-
return True
1756+
# return work_id
1757+
return result["data"]["stixCoreObjectEdit"]["askEnrichment"]["id"]

‎pycti/utils/opencti_stix2.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from pycti.utils.constants import IdentityTypes, LocationTypes, StixCyberObservableTypes
1616
from pycti.utils.opencti_stix2_splitter import OpenCTIStix2Splitter
1717
from pycti.utils.opencti_stix2_update import OpenCTIStix2Update
18+
from pycti.utils.opencti_stix2_utils import OBSERVABLES_VALUE_INT
1819

1920
datefinder.ValueError = ValueError, OverflowError
2021
utc = pytz.UTC
@@ -665,10 +666,13 @@ def import_observable(
665666
"reports": reports,
666667
}
667668
if stix_object["type"] == "x-opencti-simple-observable":
669+
print(stix_object)
668670
stix_observable_result = self.opencti.stix_cyber_observable.create(
669671
simple_observable_id=stix_object["id"],
670672
simple_observable_key=stix_object["key"],
671-
simple_observable_value=stix_object["value"],
673+
simple_observable_value=stix_object["value"]
674+
if stix_object["key"] not in OBSERVABLES_VALUE_INT
675+
else int(stix_object["value"]),
672676
simple_observable_description=stix_object["description"]
673677
if "description" in stix_object
674678
else None,

‎pycti/utils/opencti_stix2_utils.py

+6
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@
3636
"X-OpenCTI-Hostname": ["value"],
3737
}
3838

39+
OBSERVABLES_VALUE_INT = [
40+
"Autonomous-System.number",
41+
"Network-Traffic.dst_port",
42+
"Process.pid",
43+
]
44+
3945

4046
class OpenCTIStix2Utils:
4147
@staticmethod

‎tests/02-integration/stories/__init__.py ‎tests/02-integration/connector/__init__.py

File renamed without changes.

0 commit comments

Comments
 (0)