From 396c2fcac7f1a9376716d71648022dc7b570e5ab Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 20:35:02 +0000 Subject: [PATCH 01/24] al client wrapper --- .../v4_client/module/search/__init__.py | 17 +- assemblyline_client/v4_client/wrapper.py | 274 ++++++++++++++++++ 2 files changed, 289 insertions(+), 2 deletions(-) create mode 100644 assemblyline_client/v4_client/wrapper.py diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index 30a8e2d..4089c2f 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -7,7 +7,7 @@ from assemblyline_client.v4_client.module.search.histogram import Histogram from assemblyline_client.v4_client.module.search.stats import Stats from assemblyline_client.v4_client.module.search.stream import Stream - +from assemblyline_client.v4_client.wrapper import FileWrapper, AlertWrapper, BadlistWrapper, HeuristicWrapper, ResultWrapper, SignatureWrapper, SubmissionWrapper, WorkflowWrapper class Search(object): def __init__(self, connection): @@ -37,7 +37,19 @@ def _do_search(self, index, query, use_archive=False, track_total_hits=None, **k if track_total_hits: kwargs['track_total_hits'] = track_total_hits path = api_path('search', index) - return self._connection.post(path, data=json.dumps(kwargs)) + data = self._connection.post(path, data=json.dumps(kwargs)) + # wrap item dict with ALDict for added functionality + for idx, item in enumerate(data['items']): + if index == 'alert': data['items'][idx] = AlertWrapper(self._connection, item) + if index == 'badlist': data['items'][idx] = BadlistWrapper(self._connection, item) + if index == 'file': data['items'][idx] = FileWrapper(self._connection, item) + if index == 'heuristic': data['items'][idx] = HeuristicWrapper(self._connection, item) + if index == 'result': data['items'][idx] = ResultWrapper(self._connection, item) + if index == 'safelist': data['items'][idx] = SafelistWrapper(self._connection, item) + if index == 'signature': data['items'][idx] = SignatureWrapper(self._connection, item) + if index == 'submission': data['items'][idx] = SubmissionWrapper(self._connection, item) + if index == 'workflow': data['items'][idx] = WorkflowWrapper(self._connection, item) + return data def alert(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, use_archive=False, track_total_hits=None): @@ -254,3 +266,4 @@ def workflow(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, t return self._do_search('workflow', query, filters=filters, fl=fl, offset=offset, rows=rows, sort=sort, timeout=timeout, use_archive=use_archive, track_total_hits=track_total_hits) + diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py new file mode 100644 index 0000000..f228789 --- /dev/null +++ b/assemblyline_client/v4_client/wrapper.py @@ -0,0 +1,274 @@ +from assemblyline_client.v4_client.module.file import File +from assemblyline_client.v4_client.module.alert import Alert +from assemblyline_client.v4_client.module.result import Result +from assemblyline_client.v4_client.module.badlist import Badlist +from assemblyline_client.v4_client.module.safelist import Safelist +from assemblyline_client.v4_client.module.signature import Signature +from assemblyline_client.v4_client.module.submission import Submission +# AL client wrapper to allow direct assemblyline_client actions from queries + +class BaseWrapper(dict): + def __init__(self, connection, data): + self.connection = connection + super().__init__(data) + + def run(self): + print("function ran") + +class FileWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.file = File(connection) + super().__init__(connection, data) + + def download(self): + return self.file.download(self['sha256']) + + def delete_from_filestore(self): + return self.file.delete_from_filestore(self['sha256']) + + def hex(self, bytes_only=False, length=None): + return self.file.hex(self['sha256'], bytes_only, length) + + def info(self): + return self.file.info(self['sha256']) + + def result(self): + return self.file.result(self['sha256']) + + def score(self): + return self.file.score(self['sha256']) + + def strings(self): + return self.file.strings(self['sha256']) + + def children(self): + return self.file.children(self['sha256']) + + def ascii(self): + return self.file.ascii(self['sha256']) + +class AlertWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.alert = Alert(connection) + super().__init__(connection, data) + + def __call__(self): + return self.alert(self['id']) + + def grouped(self, field, fq=[], q=None, tc_start=None, tc=None, no_delay=False, offset=0, rows=10, + use_archive=False, track_total_hits=None): + return self.alert.grouped(field, fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay, offset=offset, rows=rows, use_archive=use_archive, track_total_hits=track_total_hits) + + + def label(self, *labels): + return self.alert.label(self['id'], *labels) + + def labels(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): + return self.alert.labels(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + + + def list(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False, offset=0, rows=10, + use_archive=False, track_total_hits=None): + return self.alert.list(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay, offset=offset, rows=rows, use_archive=use_archive, track_total_hits=track_total_hits) + + def ownership(self): + return self.alert.ownership(self['id']) + + def priority(self, priority): + return self.alert.priority(self['id'], priority) + + def priorities(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): + return self.alert.priorities(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + + def related(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): + return self.alert.related(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + + def remove_label(self, *labels): + return self.alert.remove_label(self['id'], *labels) + + def statistics(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): + return self.alert.statistics(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + + def status(self, status): + return self.alert.status(self['id'], status) + + def statuses(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): + return self.alert.statuses(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + + def verdict(self, verdict): + return self.alert.verdict(self['id'], verdict) + + +class BadlistWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.badlist = Badlist(connection) + super().__init__(connection, data) + + def __call__(self, qhash): + return self.badlist(qhash) + + def add_update(self, badlist_object): + return self.badlist.add_update(badlist_object) + + def add_update_many(self, list_of_badlist_object): + return self.badlist.add_update_many(list_of_badlist_object) + + def delete(self): + return self.badlist.delete(self['id']) + + def set_enabled(self, enabled): + return self.badlist.set_enabled(self['id'], enabled) + + def ssdeep(self): + return self.badlist.ssdeep(self['hashes']['ssdeep']) + + def tlsh(self): + return self.badlist.tlsh(self['hashes']['tlsh']) + + +class HeuristicWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.heuristic = Heuristic(connection) + super().__init__(connection, data) + + def __call__(self): + return self.heuristic(self['id']) + + def stats(self): + return self.heuristic.stats() + +class ResultWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.result = Result(connection) + super().__init__(connection, data) + + def __call__(self): + return self.result(self['id']) + + def error(self): + return self.result.error(self['id']) + + def multiple(self, error=None, result=None): + return self.result.multiple(error, result) + + +class SafelistWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.safelist = Safelist(connection) + super().__init__(connection, data) + + def __call__(self): + return self.safelist(self['id']) + + def add_update(self, safelist_object): + return self.safelist.add_update(safelist_object) + + def add_update_many(self, list_of_safelist_object): + return self.safelist.add_update_many(list_of_safelist_object) + + def delete(self): + return self.safelist.delete(self['id']) + + def set_enabled(self, enabled): + return self.safelist.set_enabled(self['id'],enabled) + + +class SignatureWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.signature = Signature(connection) + super().__init__(connection, data) + + def __call__(self): + return self.signature(self['id']) + + def add_update(self, data, dedup_name=True): + return self.signature.add_update(data, dedup_name=dedup_name) + + def add_update_many(self, source, sig_type, data, dedup_name=True): + return self.signature.add_update_many(source, sig_type, data, dedup_name=dedup_name) + + def change_status(self, status): + return self.signature.change_status(self['id'], status) + + def clear_status(self): + return self.signature.clear_status(self['id']) + + def delete(self): + return self.signature.delete(self['id']) + + def download(self, output=None, query=None): + return self.signature.download(output=output, query=query) + + def stats(self): + return self.signature.stats() + + def update_available(self, since='', sig_type='*'): + return self.signature.update_available(since=since, sig_type=sig_type) + +class SubmissionWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.submission = Submission(connection) + super().__init__(connection, data) + + def __call__(self): + return self.submission(self['sid']) + + def delete(self): + return self.submission.delete(self['sid']) + + def file(self, sha256, results=None, errors=None): + return self.submission.file(self['sid'], sha256, results=results, errors=errors) + + def full(self): + return self.submission.full(self['sid']) + + def is_completed(self): + return self.submission.is_completed(self['sid']) + + def list(self, user=None, group=None, fq=None, rows=10, offset=0, use_archive=False, track_total_hits=None): + return self.submission.list(user=user, group=group, fq=fq, rows=rows, offset=offset, use_archive=use_archive, track_total_hits=track_total_hits) + + def report(self): + return self.submission.report(self['sid']) + + def set_verdict(self, verdict): + return self.submission.set_verdict(self['sid'], verdict) + + def summary(self): + return self.submission.summary(self['sid']) + + def tree(self): + return self.submission.tree(self['sid']) + + +class WorkflowWrapper(BaseWrapper): + + def __init__(self, connection, data): + self.workflow = Workflow(connection) + super().__init__(connection, data) + + def __call__(self): + return self.workflow(self['workflow_id']) + + def add(self, workflow): + return self.workflow.add(workflow) + + def delete(self): + return self.workflow.delete(self['workflow_id']) + + def labels(self): + return self.workflow.labels() + + def list(self, query="*:*", rows=10, offset=0): + return self.workflow.list(query=query, rows=rows, offset=offset) + + def update(self, workflow): + return self.workflow.update(self['workflow_id'], workflow) From 60a0e2967eb168f044487a901b3a5cfdf8e496db Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 20:46:38 +0000 Subject: [PATCH 02/24] al client wrapper --- assemblyline_client/v4_client/wrapper.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index f228789..2b8de17 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -15,6 +15,7 @@ def __init__(self, connection, data): def run(self): print("function ran") + class FileWrapper(BaseWrapper): def __init__(self, connection, data): @@ -48,6 +49,7 @@ def children(self): def ascii(self): return self.file.ascii(self['sha256']) + class AlertWrapper(BaseWrapper): def __init__(self, connection, data): @@ -141,6 +143,7 @@ def __call__(self): def stats(self): return self.heuristic.stats() + class ResultWrapper(BaseWrapper): def __init__(self, connection, data): @@ -212,6 +215,7 @@ def stats(self): def update_available(self, since='', sig_type='*'): return self.signature.update_available(since=since, sig_type=sig_type) + class SubmissionWrapper(BaseWrapper): def __init__(self, connection, data): From 16181f94212efef33db63ad32b618c9cf0d4140b Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:02:35 +0000 Subject: [PATCH 03/24] bug --- assemblyline_client/v4_client/wrapper.py | 1 + 1 file changed, 1 insertion(+) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 2b8de17..217fc81 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -5,6 +5,7 @@ from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature from assemblyline_client.v4_client.module.submission import Submission +from assemblyline_client.v4_client.module.workflow import Workflow # AL client wrapper to allow direct assemblyline_client actions from queries class BaseWrapper(dict): From b62bd2915dfa07e42c5ca7ee118126ec0d87a865 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:09:12 +0000 Subject: [PATCH 04/24] bug --- assemblyline_client/v4_client/wrapper.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 217fc81..f6231fd 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -1,11 +1,12 @@ -from assemblyline_client.v4_client.module.file import File from assemblyline_client.v4_client.module.alert import Alert -from assemblyline_client.v4_client.module.result import Result from assemblyline_client.v4_client.module.badlist import Badlist +from assemblyline_client.v4_client.module.file import File +from assemblyline_client.v4_client.module.file import Heuristic +from assemblyline_client.v4_client.module.result import Result from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature from assemblyline_client.v4_client.module.submission import Submission -from assemblyline_client.v4_client.module.workflow import Workflow +from assemblyline_client.v4_client.module.workflow import Workflow: # AL client wrapper to allow direct assemblyline_client actions from queries class BaseWrapper(dict): From ebc1d36abdcb638dd86f8ff7cad9f2472d0b347b Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:12:15 +0000 Subject: [PATCH 05/24] bug --- assemblyline_client/v4_client/wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index f6231fd..8b27eea 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -6,7 +6,7 @@ from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature from assemblyline_client.v4_client.module.submission import Submission -from assemblyline_client.v4_client.module.workflow import Workflow: +from assemblyline_client.v4_client.module.workflow import Workflow # AL client wrapper to allow direct assemblyline_client actions from queries class BaseWrapper(dict): From 225116b0271fa46f4c5c21ad24ac2c80851e6090 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:15:22 +0000 Subject: [PATCH 06/24] imports --- assemblyline_client/v4_client/wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 8b27eea..90d7bbe 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -1,7 +1,7 @@ from assemblyline_client.v4_client.module.alert import Alert from assemblyline_client.v4_client.module.badlist import Badlist from assemblyline_client.v4_client.module.file import File -from assemblyline_client.v4_client.module.file import Heuristic +from assemblyline_client.v4_client.module.heuristic import Heuristic from assemblyline_client.v4_client.module.result import Result from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature From 3717bba5bc4d9e9f95eecb7f43b850216a1c568e Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:24:13 +0000 Subject: [PATCH 07/24] imports --- assemblyline_client/v4_client/wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 90d7bbe..9e5384e 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -1,7 +1,7 @@ from assemblyline_client.v4_client.module.alert import Alert from assemblyline_client.v4_client.module.badlist import Badlist from assemblyline_client.v4_client.module.file import File -from assemblyline_client.v4_client.module.heuristic import Heuristic +from assemblyline_client.v4_client.module.heuristics import Heuristic from assemblyline_client.v4_client.module.result import Result from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature From 3b2cb5b114ab42e198e5f8cbd612d054549f5d72 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:27:36 +0000 Subject: [PATCH 08/24] imports --- assemblyline_client/v4_client/wrapper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 9e5384e..6033687 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -1,7 +1,7 @@ from assemblyline_client.v4_client.module.alert import Alert from assemblyline_client.v4_client.module.badlist import Badlist from assemblyline_client.v4_client.module.file import File -from assemblyline_client.v4_client.module.heuristics import Heuristic +from assemblyline_client.v4_client.module.heuristics import Heuristics from assemblyline_client.v4_client.module.result import Result from assemblyline_client.v4_client.module.safelist import Safelist from assemblyline_client.v4_client.module.signature import Signature @@ -136,7 +136,7 @@ def tlsh(self): class HeuristicWrapper(BaseWrapper): def __init__(self, connection, data): - self.heuristic = Heuristic(connection) + self.heuristic = Heuristics(connection) super().__init__(connection, data) def __call__(self): From 2cedcdf71bf8d8b6c0a9a09d2fc71b8ca19ffe6a Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 26 Nov 2024 22:36:32 +0000 Subject: [PATCH 09/24] imports fix --- assemblyline_client/v4_client/module/search/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index 4089c2f..1bb1e6f 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -7,7 +7,7 @@ from assemblyline_client.v4_client.module.search.histogram import Histogram from assemblyline_client.v4_client.module.search.stats import Stats from assemblyline_client.v4_client.module.search.stream import Stream -from assemblyline_client.v4_client.wrapper import FileWrapper, AlertWrapper, BadlistWrapper, HeuristicWrapper, ResultWrapper, SignatureWrapper, SubmissionWrapper, WorkflowWrapper +from assemblyline_client.v4_client.wrapper import FileWrapper, AlertWrapper, BadlistWrapper, HeuristicWrapper, ResultWrapper, SafelistWrapper, SignatureWrapper, SubmissionWrapper, WorkflowWrapper class Search(object): def __init__(self, connection): From 7edeb88fb6e90adfd72690d5bfc4384f8a8b4499 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Fri, 29 Nov 2024 19:57:56 +0000 Subject: [PATCH 10/24] test --- assemblyline_client/v4_client/wrapper.py | 1 - 1 file changed, 1 deletion(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 6033687..3006568 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -7,7 +7,6 @@ from assemblyline_client.v4_client.module.signature import Signature from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow -# AL client wrapper to allow direct assemblyline_client actions from queries class BaseWrapper(dict): def __init__(self, connection, data): From ce8d22692c9a1e7eca77d85ff9c3edfe4075bfca Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 3 Dec 2024 20:18:54 +0000 Subject: [PATCH 11/24] file submission wrapper arg removed --- assemblyline_client/v4_client/wrapper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 3006568..e79c4b1 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -229,8 +229,8 @@ def __call__(self): def delete(self): return self.submission.delete(self['sid']) - def file(self, sha256, results=None, errors=None): - return self.submission.file(self['sid'], sha256, results=results, errors=errors) + def file(self, results=None, errors=None): + return self.submission.file(self['sid'], self.full()['files'][0]['sha256'], results=results, errors=errors) def full(self): return self.submission.full(self['sid']) From a6d5d64d3f44a2c3795bdf07c1f7d579271faadb Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 3 Dec 2024 22:13:44 +0000 Subject: [PATCH 12/24] wrapper cleanup --- assemblyline_client/v4_client/wrapper.py | 25 ++++++++++-------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index e79c4b1..b008cb3 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -9,19 +9,14 @@ from assemblyline_client.v4_client.module.workflow import Workflow class BaseWrapper(dict): - def __init__(self, connection, data): - self.connection = connection + def __init__(self, data): super().__init__(data) - def run(self): - print("function ran") - - class FileWrapper(BaseWrapper): def __init__(self, connection, data): self.file = File(connection) - super().__init__(connection, data) + super().__init__(data) def download(self): return self.file.download(self['sha256']) @@ -55,7 +50,7 @@ class AlertWrapper(BaseWrapper): def __init__(self, connection, data): self.alert = Alert(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.alert(self['id']) @@ -108,7 +103,7 @@ class BadlistWrapper(BaseWrapper): def __init__(self, connection, data): self.badlist = Badlist(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self, qhash): return self.badlist(qhash) @@ -136,7 +131,7 @@ class HeuristicWrapper(BaseWrapper): def __init__(self, connection, data): self.heuristic = Heuristics(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.heuristic(self['id']) @@ -149,7 +144,7 @@ class ResultWrapper(BaseWrapper): def __init__(self, connection, data): self.result = Result(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.result(self['id']) @@ -165,7 +160,7 @@ class SafelistWrapper(BaseWrapper): def __init__(self, connection, data): self.safelist = Safelist(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.safelist(self['id']) @@ -187,7 +182,7 @@ class SignatureWrapper(BaseWrapper): def __init__(self, connection, data): self.signature = Signature(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.signature(self['id']) @@ -221,7 +216,7 @@ class SubmissionWrapper(BaseWrapper): def __init__(self, connection, data): self.submission = Submission(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.submission(self['sid']) @@ -258,7 +253,7 @@ class WorkflowWrapper(BaseWrapper): def __init__(self, connection, data): self.workflow = Workflow(connection) - super().__init__(connection, data) + super().__init__(data) def __call__(self): return self.workflow(self['workflow_id']) From 2b77577ca4b7133cd5ec6e7f38daa00c4d7d6452 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Fri, 6 Dec 2024 16:44:19 +0000 Subject: [PATCH 13/24] wrapper mapping --- .../v4_client/module/search/__init__.py | 16 ++++------------ assemblyline_client/v4_client/wrapper.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index 1bb1e6f..bfb415d 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -7,7 +7,7 @@ from assemblyline_client.v4_client.module.search.histogram import Histogram from assemblyline_client.v4_client.module.search.stats import Stats from assemblyline_client.v4_client.module.search.stream import Stream -from assemblyline_client.v4_client.wrapper import FileWrapper, AlertWrapper, BadlistWrapper, HeuristicWrapper, ResultWrapper, SafelistWrapper, SignatureWrapper, SubmissionWrapper, WorkflowWrapper +from assemblyline_client.v4_client.wrapper import wrapper_map class Search(object): def __init__(self, connection): @@ -38,17 +38,9 @@ def _do_search(self, index, query, use_archive=False, track_total_hits=None, **k kwargs['track_total_hits'] = track_total_hits path = api_path('search', index) data = self._connection.post(path, data=json.dumps(kwargs)) - # wrap item dict with ALDict for added functionality - for idx, item in enumerate(data['items']): - if index == 'alert': data['items'][idx] = AlertWrapper(self._connection, item) - if index == 'badlist': data['items'][idx] = BadlistWrapper(self._connection, item) - if index == 'file': data['items'][idx] = FileWrapper(self._connection, item) - if index == 'heuristic': data['items'][idx] = HeuristicWrapper(self._connection, item) - if index == 'result': data['items'][idx] = ResultWrapper(self._connection, item) - if index == 'safelist': data['items'][idx] = SafelistWrapper(self._connection, item) - if index == 'signature': data['items'][idx] = SignatureWrapper(self._connection, item) - if index == 'submission': data['items'][idx] = SubmissionWrapper(self._connection, item) - if index == 'workflow': data['items'][idx] = WorkflowWrapper(self._connection, item) + + data['items'] = [wrapper_map[index](self._connection, item) for item in data['items']] + return data def alert(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index b008cb3..43e25c6 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -272,3 +272,18 @@ def list(self, query="*:*", rows=10, offset=0): def update(self, workflow): return self.workflow.update(self['workflow_id'], workflow) + + + +wrapper_map = { + 'file': FileWrapper, + 'alert': AlertWrapper, + 'badlist': BadlistWrapper, + 'heuristic': HeuristicWrapper, + 'result': ResultWrapper, + 'safelist': SafelistWrapper, + 'signature': SignatureWrapper, + 'submission': SubmissionWrapper, + 'workflow': WorkflowWrapper +} + From 82ba5d7a59585c0e44cef6bc0ae0aee2ecd832ce Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Mon, 16 Dec 2024 20:30:15 +0000 Subject: [PATCH 14/24] updated client wrapper --- .../v4_client/module/search/__init__.py | 6 +- assemblyline_client/v4_client/wrapper.py | 450 ++++++++++-------- 2 files changed, 263 insertions(+), 193 deletions(-) diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index bfb415d..34cb3fe 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -7,7 +7,7 @@ from assemblyline_client.v4_client.module.search.histogram import Histogram from assemblyline_client.v4_client.module.search.stats import Stats from assemblyline_client.v4_client.module.search.stream import Stream -from assemblyline_client.v4_client.wrapper import wrapper_map +from assemblyline_client.v4_client.wrapper import wrapper_map, BaseWrapper class Search(object): def __init__(self, connection): @@ -38,8 +38,8 @@ def _do_search(self, index, query, use_archive=False, track_total_hits=None, **k kwargs['track_total_hits'] = track_total_hits path = api_path('search', index) data = self._connection.post(path, data=json.dumps(kwargs)) - - data['items'] = [wrapper_map[index](self._connection, item) for item in data['items']] + wrapper = wrapper_map.get(index, BaseWrapper) + data['items'] = [wrapper(self, item) for item in data['items']] return data diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 43e25c6..8359c2c 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,270 +8,340 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow -class BaseWrapper(dict): - def __init__(self, data): - super().__init__(data) - -class FileWrapper(BaseWrapper): +import pandas + +wrapper_search = None + +def wrapper_function(inner_func,args=[]): + """ +Decorator for wrapper functions to provide docstrings for documentation. + +Required: +inner_func: Function that is being wrapped (.ie function called inside of a wrapper function) +args: List of required arguments to not remove from an inner function's docstring + +Sets the __doc__ attribute of the decorated function to the modified __doc__ of the inner function. + """ + + def decorator(func): + hits = ['Required:', 'Optional:'] + clear = False + docstring = [] + for line in inner_func.__doc__.splitlines(): + if line in hits: + if len(args) != 0: docstring.append(line) + clear = True + if line == '': + clear = False + if not clear or line.split(":")[0].strip() in args: + docstring.append(line) + func.__doc__ = "\n".join(docstring) + return func + + return decorator - def __init__(self, connection, data): - self.file = File(connection) +class BaseWrapper(dict): + def __init__(self, search, data): + self.search = search super().__init__(data) - def download(self): - return self.file.download(self['sha256']) - - def delete_from_filestore(self): - return self.file.delete_from_filestore(self['sha256']) + def to_dataframe(self): + return pandas.DataFrame(self['items']) - def hex(self, bytes_only=False, length=None): - return self.file.hex(self['sha256'], bytes_only, length) - - def info(self): - return self.file.info(self['sha256']) - - def result(self): - return self.file.result(self['sha256']) - - def score(self): - return self.file.score(self['sha256']) - - def strings(self): - return self.file.strings(self['sha256']) +class FileWrapper(BaseWrapper): - def children(self): - return self.file.children(self['sha256']) + def __init__(self, search, data): + self.file = File(search._connection) + super().__init__(search, data) + + @wrapper_function(File.download) + def download(self, *args, **kwargs): + return self.file.download(self['sha256'], *args, **kwargs) + + @wrapper_function(File.delete_from_filestore) + def delete_from_filestore(self, *args, **kwargs): + return self.file.delete_from_filestore(self['sha256'], *args, **kwargs) + + @wrapper_function(File.hex) + def hex(self, *args, **kwargs): + return self.file.hex(self['sha256'], *args, **kwargs) - def ascii(self): - return self.file.ascii(self['sha256']) + @wrapper_function(File.info) + def info(self, *args, **kwargs): + return self.file.info(self['sha256'], *args, **kwargs) + + @wrapper_function(File.result) + def result(self, *args, **kwargs): + return self.file.result(self['sha256'], *args, **kwargs) + @wrapper_function(File.score) + def score(self, *args, **kwargs): + return self.file.score(self['sha256'], *args, **kwargs) + + @wrapper_function(File.strings) + def strings(self, *args, **kwargs): + return self.file.strings(self['sha256'], *args, **kwargs) + + @wrapper_function(File.children) + def children(self, *args, **kawrgs): + return self.file.children(self['sha256'], *args, **kwargs) + + @wrapper_function(File.ascii) + def ascii(self, *args, **kwargs): + return self.file.ascii(self['sha256'], *args, **kwargs) + + def get_submissions(self): + """\ +Get the submissions from a file + +Returns a list of SubmissionWrapper +""" + return self.search.submission(self['sha256'])['items'] + + def get_results(self): + """\ +Get the results of a file + +Returns a list of ResultWrapper +""" + return self.search.result(f"sha256:{self['sha256']}")['items'] + + def get_extracted_files(self): + """\ +Get extracted files from a file. This searches for all results related to a sha256. + +Returns a list of FileWrapper +""" + results = self.get_results() + try: + extracted_files = [result()['response']['extracted'] for result in self.get_results()] + + query = "" + for files in extracted_files: + if len(files) == 0: continue + for idx, file in enumerate(files): + if idx+1 == len(files): + query += f"{file['sha256']}" + continue + query += f"{file['sha256']} OR " + return self.search.file(query)['items'] + + except KeyError: + return [] class AlertWrapper(BaseWrapper): - def __init__(self, connection, data): - self.alert = Alert(connection) - super().__init__(data) - + def __init__(self, search, data): + self.alert = Alert(search._connection) + super().__init__(search, data) + + @wrapper_function(Alert.__call__) def __call__(self): return self.alert(self['id']) - def grouped(self, field, fq=[], q=None, tc_start=None, tc=None, no_delay=False, offset=0, rows=10, - use_archive=False, track_total_hits=None): - return self.alert.grouped(field, fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay, offset=offset, rows=rows, use_archive=use_archive, track_total_hits=track_total_hits) - - - def label(self, *labels): - return self.alert.label(self['id'], *labels) - - def labels(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): - return self.alert.labels(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) - - - def list(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False, offset=0, rows=10, - use_archive=False, track_total_hits=None): - return self.alert.list(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay, offset=offset, rows=rows, use_archive=use_archive, track_total_hits=track_total_hits) + @wrapper_function(Alert.label, ["*labels"]) + def label(self, *args, **kwargs): + return self.alert.label(self['id'], *args, **kwargs) + @wrapper_function(Alert.ownership) def ownership(self): return self.alert.ownership(self['id']) - def priority(self, priority): - return self.alert.priority(self['id'], priority) - - def priorities(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): - return self.alert.priorities(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) - - def related(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): - return self.alert.related(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + @wrapper_function(Alert.priority, ["priority"]) + def priority(self, *args, **kwargs): + return self.alert.priority(self['id'], *args, **kwargs) - def remove_label(self, *labels): - return self.alert.remove_label(self['id'], *labels) - - def statistics(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): - return self.alert.statistics(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) - - def status(self, status): - return self.alert.status(self['id'], status) - - def statuses(self, fq=[], q=None, tc_start=None, tc=None, no_delay=False): - return self.alert.statuses(fq=fq, q=q, tc_start=tc_start, tc=tc, no_delay=no_delay) + @wrapper_function(Alert.remove_label, ["*labels"]) + def remove_label(self, *args, **kwargs): + return self.alert.remove_label(self['id'], *args, **kwargs) + + @wrapper_function(Alert.status, ["status"]) + def status(self, *args, **kwargs): + return self.alert.status(self['id'], *args, **kwargs) - def verdict(self, verdict): - return self.alert.verdict(self['id'], verdict) + @wrapper_function(Alert.verdict, ["verdict"]) + def verdict(self, *args, **kwargs): + return self.alert.verdict(self['id'], *args, **kwargs) class BadlistWrapper(BaseWrapper): - def __init__(self, connection, data): - self.badlist = Badlist(connection) - super().__init__(data) - - def __call__(self, qhash): - return self.badlist(qhash) - - def add_update(self, badlist_object): - return self.badlist.add_update(badlist_object) - - def add_update_many(self, list_of_badlist_object): - return self.badlist.add_update_many(list_of_badlist_object) - - def delete(self): - return self.badlist.delete(self['id']) - - def set_enabled(self, enabled): - return self.badlist.set_enabled(self['id'], enabled) + def __init__(self, search, data): + self.badlist = Badlist(search._connection) + super().__init__(search, data) + + @wrapper_function(Badlist.__call__) + def __call__(self, *args, **kwargs): + return self.badlist(*args, **kwargs) - def ssdeep(self): - return self.badlist.ssdeep(self['hashes']['ssdeep']) + @wrapper_function(Badlist.delete) + def delete(self, *args, **kwargs): + return self.badlist.delete(self['id'], *args, **kwargs) + + @wrapper_function(Badlist.set_enabled, ["enabled"]) + def set_enabled(self, *args, **kwargs): + return self.badlist.set_enabled(self['id'], *args, **kwargs) - def tlsh(self): - return self.badlist.tlsh(self['hashes']['tlsh']) + @wrapper_function(Badlist.ssdeep) + def ssdeep(self, *args, **kwargs): + return self.badlist.ssdeep(self['hashes']['ssdeep'], *args, **kwargs) + + @wrapper_function(Badlist.tlsh) + def tlsh(self, *args, **kwargs): + return self.badlist.tlsh(self['hashes']['tlsh'], *args, **kwargs) class HeuristicWrapper(BaseWrapper): - def __init__(self, connection, data): - self.heuristic = Heuristics(connection) - super().__init__(data) - + def __init__(self, search, data): + self.heuristic = Heuristics(search._connection) + super().__init__(search, data) + + @wrapper_function(Heuristics.__call__) def __call__(self): return self.heuristic(self['id']) - def stats(self): - return self.heuristic.stats() + @wrapper_function(Heuristics.stats) + def stats(self, *args, **kwargs): + return self.heuristic.stats(*args, **kwargs) class ResultWrapper(BaseWrapper): - def __init__(self, connection, data): - self.result = Result(connection) - super().__init__(data) + def __init__(self, search, data): + self.result = Result(search._connection) + super().__init__(search, data) + @wrapper_function(Result.__call__) def __call__(self): return self.result(self['id']) - def error(self): - return self.result.error(self['id']) - - def multiple(self, error=None, result=None): - return self.result.multiple(error, result) - + @wrapper_function(Result.error) + def error(self, *args, **kwargs): + return self.result.error(self['id'], *args, **kwargs) + + #def get_extracted_files(self): + # return self['response'] class SafelistWrapper(BaseWrapper): - def __init__(self, connection, data): - self.safelist = Safelist(connection) - super().__init__(data) - + def __init__(self, search, data): + self.safelist = Safelist(search._connection) + super().__init__(search, data) + + @wrapper_function(Safelist.__call__) def __call__(self): return self.safelist(self['id']) + + @wrapper_function(Safelist.delete) + def delete(self, *args, **kwargs): + return self.safelist.delete(self['id'], *args, **kwargs) - def add_update(self, safelist_object): - return self.safelist.add_update(safelist_object) - - def add_update_many(self, list_of_safelist_object): - return self.safelist.add_update_many(list_of_safelist_object) - - def delete(self): - return self.safelist.delete(self['id']) - - def set_enabled(self, enabled): - return self.safelist.set_enabled(self['id'],enabled) + @wrapper_function(Safelist.set_enabled, ["enabled"]) + def set_enabled(self, *args, **kwargs): + return self.safelist.set_enabled(self['id'], *args, **kwargs) class SignatureWrapper(BaseWrapper): - def __init__(self, connection, data): - self.signature = Signature(connection) - super().__init__(data) - + def __init__(self, search, data): + self.signature = Signature(search._connection) + super().__init__(search, data) + + @wrapper_function(Signature.__call__) def __call__(self): return self.signature(self['id']) + + @wrapper_function(Signature.change_status, ["status"]) + def change_status(self, *args, **kwargs): + return self.signature.change_status(self['id'], *args, **kwargs) - def add_update(self, data, dedup_name=True): - return self.signature.add_update(data, dedup_name=dedup_name) - - def add_update_many(self, source, sig_type, data, dedup_name=True): - return self.signature.add_update_many(source, sig_type, data, dedup_name=dedup_name) - - def change_status(self, status): - return self.signature.change_status(self['id'], status) - - def clear_status(self): - return self.signature.clear_status(self['id']) - - def delete(self): - return self.signature.delete(self['id']) - - def download(self, output=None, query=None): - return self.signature.download(output=output, query=query) - - def stats(self): - return self.signature.stats() + @wrapper_function(Signature.clear_status) + def clear_status(self, *args, **kwargs): + return self.signature.clear_status(self['id'], *args, **kwargs) - def update_available(self, since='', sig_type='*'): - return self.signature.update_available(since=since, sig_type=sig_type) + @wrapper_function(Signature.delete) + def delete(self, *args, **kwargs): + return self.signature.delete(self['id'], *args, **kwargs) class SubmissionWrapper(BaseWrapper): - def __init__(self, connection, data): - self.submission = Submission(connection) - super().__init__(data) - + def __init__(self, search, data): + self.submission = Submission(search._connection) + super().__init__(search, data) + + @wrapper_function(Submission.__call__) def __call__(self): return self.submission(self['sid']) - def delete(self): - return self.submission.delete(self['sid']) - - def file(self, results=None, errors=None): - return self.submission.file(self['sid'], self.full()['files'][0]['sha256'], results=results, errors=errors) - - def full(self): - return self.submission.full(self['sid']) - - def is_completed(self): - return self.submission.is_completed(self['sid']) + @wrapper_function(Submission.delete) + def delete(self, *args, **kwargs): + return self.submission.delete(self['sid'], *args, **kwargs) - def list(self, user=None, group=None, fq=None, rows=10, offset=0, use_archive=False, track_total_hits=None): - return self.submission.list(user=user, group=group, fq=fq, rows=rows, offset=offset, use_archive=use_archive, track_total_hits=track_total_hits) - - def report(self): - return self.submission.report(self['sid']) - - def set_verdict(self, verdict): - return self.submission.set_verdict(self['sid'], verdict) + @wrapper_function(Submission.file, ["sha256"]) + def file(self, *args, **kwargs): + return self.submission.file(self['sid'], *args, **kwargs) + + @wrapper_function(Submission.full) + def full(self, *args, **kwargs): + return self.submission.full(self['sid'], *args, **kwargs) - def summary(self): - return self.submission.summary(self['sid']) + @wrapper_function(Submission.is_completed) + def is_completed(self, *args, **kwargs): + return self.submission.is_completed(self['sid'], *args, **kwargs) - def tree(self): - return self.submission.tree(self['sid']) + @wrapper_function(Submission.report) + def report(self, *args, **kwargs): + return self.submission.report(self['sid'], *args, **kwargs) + @wrapper_function(Submission.set_verdict, ["verdict"]) + def set_verdict(self, *args, **kwargs): + return self.submission.set_verdict(self['sid'], *args, **kwargs) + + @wrapper_function(Submission.summary) + def summary(self, *args, **kwargs): + return self.submission.summary(self['sid'], *args, **kwargs) + @wrapper_function(Submission.tree) + def tree(self, *args, **kwargs): + return self.submission.tree(self['sid'], *args, **kwargs) + + def get_files(self): + """\ +Get the list of files that were originally submitted + +Returns a list of FileWrapper +""" + try: + submission_files = self.full()['files'] + query = "" + for idx, file in enumerate(submission_files): + if idx+1 == len(submission_files): + query += f"{file['sha256']}" + continue + query += f"{file['sha256']} OR " + + return self.search.file(query)['items'] + except KeyError: + return [] class WorkflowWrapper(BaseWrapper): - def __init__(self, connection, data): - self.workflow = Workflow(connection) - super().__init__(data) + def __init__(self, search, data): + self.workflow = Workflow(search._connection) + super().__init__(search, data) + @wrapper_function(Workflow.__call__) def __call__(self): return self.workflow(self['workflow_id']) - - def add(self, workflow): - return self.workflow.add(workflow) - + + @wrapper_function(Workflow.delete) def delete(self): return self.workflow.delete(self['workflow_id']) - def labels(self): - return self.workflow.labels() - - def list(self, query="*:*", rows=10, offset=0): - return self.workflow.list(query=query, rows=rows, offset=offset) - - def update(self, workflow): - return self.workflow.update(self['workflow_id'], workflow) + @wrapper_function(Workflow.update, ["workflow"]) + def update(self, *args, **kwargs): + return self.workflow.update(self['workflow_id'], *args, **kwargs) From e512f80f5286b12ab7d3ead0a474ee24a76e7fee Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 17:24:12 +0000 Subject: [PATCH 15/24] pandas removed --- assemblyline_client/v4_client/wrapper.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 8359c2c..61c8287 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -92,17 +92,17 @@ def ascii(self, *args, **kwargs): def get_submissions(self): """\ -Get the submissions from a file +Get the submissions from a file. -Returns a list of SubmissionWrapper +Returns a list of SubmissionWrapper. """ return self.search.submission(self['sha256'])['items'] def get_results(self): """\ -Get the results of a file +Get the results of a file. -Returns a list of ResultWrapper +Returns a list of ResultWrapper. """ return self.search.result(f"sha256:{self['sha256']}")['items'] @@ -110,7 +110,7 @@ def get_extracted_files(self): """\ Get extracted files from a file. This searches for all results related to a sha256. -Returns a list of FileWrapper +Returns a list of FileWrapper. """ results = self.get_results() try: From 1a0485decee72f8cbbf643777a1d3a2a43e17812 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 18:28:20 +0000 Subject: [PATCH 16/24] new line --- assemblyline_client/v4_client/module/search/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index 34cb3fe..c8a5974 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -258,4 +258,3 @@ def workflow(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, t return self._do_search('workflow', query, filters=filters, fl=fl, offset=offset, rows=rows, sort=sort, timeout=timeout, use_archive=use_archive, track_total_hits=track_total_hits) - From 01f53b30458ebea44c9585aff31b1ef201ea2937 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 19:44:57 +0000 Subject: [PATCH 17/24] removed funcs --- assemblyline_client/v4_client/wrapper.py | 87 +++++++++++------------- 1 file changed, 39 insertions(+), 48 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 61c8287..c130da1 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,15 +8,14 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow -import pandas - wrapper_search = None -def wrapper_function(inner_func,args=[]): + +def wrapper_function(inner_func, args=[]): """ Decorator for wrapper functions to provide docstrings for documentation. -Required: +Required: inner_func: Function that is being wrapped (.ie function called inside of a wrapper function) args: List of required arguments to not remove from an inner function's docstring @@ -29,7 +28,8 @@ def decorator(func): docstring = [] for line in inner_func.__doc__.splitlines(): if line in hits: - if len(args) != 0: docstring.append(line) + if len(args) != 0: + docstring.append(line) clear = True if line == '': clear = False @@ -40,28 +40,27 @@ def decorator(func): return decorator + class BaseWrapper(dict): def __init__(self, search, data): self.search = search super().__init__(data) - def to_dataframe(self): - return pandas.DataFrame(self['items']) class FileWrapper(BaseWrapper): def __init__(self, search, data): self.file = File(search._connection) super().__init__(search, data) - + @wrapper_function(File.download) def download(self, *args, **kwargs): return self.file.download(self['sha256'], *args, **kwargs) - + @wrapper_function(File.delete_from_filestore) def delete_from_filestore(self, *args, **kwargs): return self.file.delete_from_filestore(self['sha256'], *args, **kwargs) - + @wrapper_function(File.hex) def hex(self, *args, **kwargs): return self.file.hex(self['sha256'], *args, **kwargs) @@ -69,7 +68,7 @@ def hex(self, *args, **kwargs): @wrapper_function(File.info) def info(self, *args, **kwargs): return self.file.info(self['sha256'], *args, **kwargs) - + @wrapper_function(File.result) def result(self, *args, **kwargs): return self.file.result(self['sha256'], *args, **kwargs) @@ -77,15 +76,15 @@ def result(self, *args, **kwargs): @wrapper_function(File.score) def score(self, *args, **kwargs): return self.file.score(self['sha256'], *args, **kwargs) - + @wrapper_function(File.strings) def strings(self, *args, **kwargs): return self.file.strings(self['sha256'], *args, **kwargs) - + @wrapper_function(File.children) def children(self, *args, **kawrgs): return self.file.children(self['sha256'], *args, **kwargs) - + @wrapper_function(File.ascii) def ascii(self, *args, **kwargs): return self.file.ascii(self['sha256'], *args, **kwargs) @@ -112,15 +111,16 @@ def get_extracted_files(self): Returns a list of FileWrapper. """ - results = self.get_results() try: - extracted_files = [result()['response']['extracted'] for result in self.get_results()] + extracted_files = [result()['response']['extracted'] + for result in self.get_results()] query = "" for files in extracted_files: - if len(files) == 0: continue + if len(files) == 0: + continue for idx, file in enumerate(files): - if idx+1 == len(files): + if idx + 1 == len(files): query += f"{file['sha256']}" continue query += f"{file['sha256']} OR " @@ -129,12 +129,13 @@ def get_extracted_files(self): except KeyError: return [] + class AlertWrapper(BaseWrapper): def __init__(self, search, data): self.alert = Alert(search._connection) - super().__init__(search, data) - + super().__init__(search, data) + @wrapper_function(Alert.__call__) def __call__(self): return self.alert(self['id']) @@ -154,7 +155,7 @@ def priority(self, *args, **kwargs): @wrapper_function(Alert.remove_label, ["*labels"]) def remove_label(self, *args, **kwargs): return self.alert.remove_label(self['id'], *args, **kwargs) - + @wrapper_function(Alert.status, ["status"]) def status(self, *args, **kwargs): return self.alert.status(self['id'], *args, **kwargs) @@ -165,11 +166,11 @@ def verdict(self, *args, **kwargs): class BadlistWrapper(BaseWrapper): - + def __init__(self, search, data): self.badlist = Badlist(search._connection) super().__init__(search, data) - + @wrapper_function(Badlist.__call__) def __call__(self, *args, **kwargs): return self.badlist(*args, **kwargs) @@ -177,26 +178,18 @@ def __call__(self, *args, **kwargs): @wrapper_function(Badlist.delete) def delete(self, *args, **kwargs): return self.badlist.delete(self['id'], *args, **kwargs) - + @wrapper_function(Badlist.set_enabled, ["enabled"]) def set_enabled(self, *args, **kwargs): return self.badlist.set_enabled(self['id'], *args, **kwargs) - @wrapper_function(Badlist.ssdeep) - def ssdeep(self, *args, **kwargs): - return self.badlist.ssdeep(self['hashes']['ssdeep'], *args, **kwargs) - - @wrapper_function(Badlist.tlsh) - def tlsh(self, *args, **kwargs): - return self.badlist.tlsh(self['hashes']['tlsh'], *args, **kwargs) - class HeuristicWrapper(BaseWrapper): def __init__(self, search, data): self.heuristic = Heuristics(search._connection) super().__init__(search, data) - + @wrapper_function(Heuristics.__call__) def __call__(self): return self.heuristic(self['id']) @@ -219,20 +212,18 @@ def __call__(self): @wrapper_function(Result.error) def error(self, *args, **kwargs): return self.result.error(self['id'], *args, **kwargs) - - #def get_extracted_files(self): - # return self['response'] + class SafelistWrapper(BaseWrapper): def __init__(self, search, data): self.safelist = Safelist(search._connection) super().__init__(search, data) - + @wrapper_function(Safelist.__call__) def __call__(self): return self.safelist(self['id']) - + @wrapper_function(Safelist.delete) def delete(self, *args, **kwargs): return self.safelist.delete(self['id'], *args, **kwargs) @@ -247,11 +238,11 @@ class SignatureWrapper(BaseWrapper): def __init__(self, search, data): self.signature = Signature(search._connection) super().__init__(search, data) - + @wrapper_function(Signature.__call__) def __call__(self): return self.signature(self['id']) - + @wrapper_function(Signature.change_status, ["status"]) def change_status(self, *args, **kwargs): return self.signature.change_status(self['id'], *args, **kwargs) @@ -270,7 +261,7 @@ class SubmissionWrapper(BaseWrapper): def __init__(self, search, data): self.submission = Submission(search._connection) super().__init__(search, data) - + @wrapper_function(Submission.__call__) def __call__(self): return self.submission(self['sid']) @@ -282,7 +273,7 @@ def delete(self, *args, **kwargs): @wrapper_function(Submission.file, ["sha256"]) def file(self, *args, **kwargs): return self.submission.file(self['sid'], *args, **kwargs) - + @wrapper_function(Submission.full) def full(self, *args, **kwargs): return self.submission.full(self['sid'], *args, **kwargs) @@ -298,10 +289,11 @@ def report(self, *args, **kwargs): @wrapper_function(Submission.set_verdict, ["verdict"]) def set_verdict(self, *args, **kwargs): return self.submission.set_verdict(self['sid'], *args, **kwargs) - + @wrapper_function(Submission.summary) def summary(self, *args, **kwargs): return self.submission.summary(self['sid'], *args, **kwargs) + @wrapper_function(Submission.tree) def tree(self, *args, **kwargs): return self.submission.tree(self['sid'], *args, **kwargs) @@ -316,7 +308,7 @@ def get_files(self): submission_files = self.full()['files'] query = "" for idx, file in enumerate(submission_files): - if idx+1 == len(submission_files): + if idx + 1 == len(submission_files): query += f"{file['sha256']}" continue query += f"{file['sha256']} OR " @@ -325,8 +317,9 @@ def get_files(self): except KeyError: return [] + class WorkflowWrapper(BaseWrapper): - + def __init__(self, search, data): self.workflow = Workflow(search._connection) super().__init__(search, data) @@ -334,7 +327,7 @@ def __init__(self, search, data): @wrapper_function(Workflow.__call__) def __call__(self): return self.workflow(self['workflow_id']) - + @wrapper_function(Workflow.delete) def delete(self): return self.workflow.delete(self['workflow_id']) @@ -344,7 +337,6 @@ def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) - wrapper_map = { 'file': FileWrapper, 'alert': AlertWrapper, @@ -356,4 +348,3 @@ def update(self, *args, **kwargs): 'submission': SubmissionWrapper, 'workflow': WorkflowWrapper } - From 5c1a4d0f68951320fe572f88f4d1885683bd0063 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 19:48:39 +0000 Subject: [PATCH 18/24] removed stats --- assemblyline_client/v4_client/wrapper.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index c130da1..d73e95a 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -194,10 +194,6 @@ def __init__(self, search, data): def __call__(self): return self.heuristic(self['id']) - @wrapper_function(Heuristics.stats) - def stats(self, *args, **kwargs): - return self.heuristic.stats(*args, **kwargs) - class ResultWrapper(BaseWrapper): From 7ea3d9ead320439078bf6d2c27c8e2355f22293e Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 21:13:58 +0000 Subject: [PATCH 19/24] fix --- assemblyline_client/v4_client/wrapper.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index d73e95a..9aa4fa1 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,11 +8,8 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow -wrapper_search = None - - def wrapper_function(inner_func, args=[]): - """ +"""\ Decorator for wrapper functions to provide docstrings for documentation. Required: @@ -20,7 +17,7 @@ def wrapper_function(inner_func, args=[]): args: List of required arguments to not remove from an inner function's docstring Sets the __doc__ attribute of the decorated function to the modified __doc__ of the inner function. - """ +""" def decorator(func): hits = ['Required:', 'Optional:'] @@ -82,7 +79,7 @@ def strings(self, *args, **kwargs): return self.file.strings(self['sha256'], *args, **kwargs) @wrapper_function(File.children) - def children(self, *args, **kawrgs): + def children(self, *args, **kwargs): return self.file.children(self['sha256'], *args, **kwargs) @wrapper_function(File.ascii) From ee1f7d104737e36700c47af6709c4c8d6a489161 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Tue, 17 Dec 2024 21:49:41 +0000 Subject: [PATCH 20/24] wrapper global var renamed --- assemblyline_client/v4_client/wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 9aa4fa1..0db3fe7 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -330,7 +330,7 @@ def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) -wrapper_map = { +WRAPPER_MAP = { 'file': FileWrapper, 'alert': AlertWrapper, 'badlist': BadlistWrapper, From d10368a24513716235b803ccd40229d163549264 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Wed, 18 Dec 2024 17:19:21 +0000 Subject: [PATCH 21/24] global var --- assemblyline_client/v4_client/wrapper.py | 25 ++++++++++++------------ 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 0db3fe7..a5ded1f 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,6 +8,18 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow +WRAPPER_MAP = { + 'file': FileWrapper, + 'alert': AlertWrapper, + 'badlist': BadlistWrapper, + 'heuristic': HeuristicWrapper, + 'result': ResultWrapper, + 'safelist': SafelistWrapper, + 'signature': SignatureWrapper, + 'submission': SubmissionWrapper, + 'workflow': WorkflowWrapper +} + def wrapper_function(inner_func, args=[]): """\ Decorator for wrapper functions to provide docstrings for documentation. @@ -328,16 +340,3 @@ def delete(self): @wrapper_function(Workflow.update, ["workflow"]) def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) - - -WRAPPER_MAP = { - 'file': FileWrapper, - 'alert': AlertWrapper, - 'badlist': BadlistWrapper, - 'heuristic': HeuristicWrapper, - 'result': ResultWrapper, - 'safelist': SafelistWrapper, - 'signature': SignatureWrapper, - 'submission': SubmissionWrapper, - 'workflow': WorkflowWrapper -} From 1b20c0a8c0a3fc929cf8570de84f6876b6c9b6f2 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Thu, 19 Dec 2024 16:16:53 +0000 Subject: [PATCH 22/24] pipeline fix --- .../v4_client/module/search/__init__.py | 272 ++++++++++++++---- assemblyline_client/v4_client/wrapper.py | 26 +- 2 files changed, 229 insertions(+), 69 deletions(-) diff --git a/assemblyline_client/v4_client/module/search/__init__.py b/assemblyline_client/v4_client/module/search/__init__.py index c8a5974..a89eb61 100644 --- a/assemblyline_client/v4_client/module/search/__init__.py +++ b/assemblyline_client/v4_client/module/search/__init__.py @@ -7,7 +7,8 @@ from assemblyline_client.v4_client.module.search.histogram import Histogram from assemblyline_client.v4_client.module.search.stats import Stats from assemblyline_client.v4_client.module.search.stream import Stream -from assemblyline_client.v4_client.wrapper import wrapper_map, BaseWrapper +from assemblyline_client.v4_client.wrapper import WRAPPER_MAP, BaseWrapper + class Search(object): def __init__(self, connection): @@ -19,7 +20,13 @@ def __init__(self, connection): self.stats = Stats(connection) self.stream = Stream(connection, self._do_search) - def _do_search(self, index, query, use_archive=False, track_total_hits=None, **kwargs): + def _do_search( + self, + index, + query, + use_archive=False, + track_total_hits=None, + **kwargs): if index not in SEARCHABLE: raise ClientError("Index %s is not searchable" % index, 400) @@ -38,13 +45,22 @@ def _do_search(self, index, query, use_archive=False, track_total_hits=None, **k kwargs['track_total_hits'] = track_total_hits path = api_path('search', index) data = self._connection.post(path, data=json.dumps(kwargs)) - wrapper = wrapper_map.get(index, BaseWrapper) + wrapper = WRAPPER_MAP.get(index, BaseWrapper) data['items'] = [wrapper(self, item) for item in data['items']] return data - def alert(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + def alert( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search alerts with a lucene query. @@ -63,12 +79,29 @@ def alert(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, time Returns all results. """ - return self._do_search('alert', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def badlist(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'alert', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def badlist( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search badlists with a lucene query. @@ -87,12 +120,29 @@ def badlist(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, ti Returns all results. """ - return self._do_search('badlist', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def file(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'badlist', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def file( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search files with a lucene query. @@ -111,12 +161,29 @@ def file(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeo Returns all results. """ - return self._do_search('file', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def heuristic(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'file', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def heuristic( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search heuristics with a lucene query. @@ -135,12 +202,29 @@ def heuristic(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, Returns all results. """ - return self._do_search('heuristic', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def result(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'heuristic', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def result( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search results with a lucene query. @@ -159,12 +243,29 @@ def result(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, tim Returns all results. """ - return self._do_search('result', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def safelist(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'result', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def safelist( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search safelist with a lucene query. @@ -183,12 +284,29 @@ def safelist(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, t Returns all results. """ - return self._do_search('safelist', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def signature(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'safelist', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def signature( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search signatures with a lucene query. @@ -207,12 +325,29 @@ def signature(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, Returns all results. """ - return self._do_search('signature', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def submission(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'signature', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def submission( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search submissions with a lucene query. @@ -231,12 +366,29 @@ def submission(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, Returns all results. """ - return self._do_search('submission', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) - - def workflow(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None, - use_archive=False, track_total_hits=None): + return self._do_search( + 'submission', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) + + def workflow( + self, + query, + filters=None, + fl=None, + offset=0, + rows=25, + sort=None, + timeout=None, + use_archive=False, + track_total_hits=None): """\ Search workflow with a lucene query. @@ -255,6 +407,14 @@ def workflow(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, t Returns all results. """ - return self._do_search('workflow', query, filters=filters, fl=fl, offset=offset, - rows=rows, sort=sort, timeout=timeout, - use_archive=use_archive, track_total_hits=track_total_hits) + return self._do_search( + 'workflow', + query, + filters=filters, + fl=fl, + offset=offset, + rows=rows, + sort=sort, + timeout=timeout, + use_archive=use_archive, + track_total_hits=track_total_hits) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index a5ded1f..f45cc2b 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,20 +8,8 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow -WRAPPER_MAP = { - 'file': FileWrapper, - 'alert': AlertWrapper, - 'badlist': BadlistWrapper, - 'heuristic': HeuristicWrapper, - 'result': ResultWrapper, - 'safelist': SafelistWrapper, - 'signature': SignatureWrapper, - 'submission': SubmissionWrapper, - 'workflow': WorkflowWrapper -} - def wrapper_function(inner_func, args=[]): -"""\ + """\ Decorator for wrapper functions to provide docstrings for documentation. Required: @@ -340,3 +328,15 @@ def delete(self): @wrapper_function(Workflow.update, ["workflow"]) def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) + +WRAPPER_MAP = { + 'file': FileWrapper, + 'alert': AlertWrapper, + 'badlist': BadlistWrapper, + 'heuristic': HeuristicWrapper, + 'result': ResultWrapper, + 'safelist': SafelistWrapper, + 'signature': SignatureWrapper, + 'submission': SubmissionWrapper, + 'workflow': WorkflowWrapper +} From 24be2e94ed32629f83d9e053fc0833a0bc415a44 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Thu, 19 Dec 2024 16:30:11 +0000 Subject: [PATCH 23/24] pipeline fix --- assemblyline_client/v4_client/wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index f45cc2b..03a3166 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,6 +8,7 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow + def wrapper_function(inner_func, args=[]): """\ Decorator for wrapper functions to provide docstrings for documentation. @@ -329,6 +330,7 @@ def delete(self): def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) + WRAPPER_MAP = { 'file': FileWrapper, 'alert': AlertWrapper, From 87ac97724185c4251ff3bfe7cb2052da398a7a08 Mon Sep 17 00:00:00 2001 From: cccs-ach Date: Thu, 19 Dec 2024 19:56:29 +0000 Subject: [PATCH 24/24] wrapper tests --- assemblyline_client/v4_client/wrapper.py | 5 ++- test/test_wrapper.py | 43 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 3 deletions(-) create mode 100644 test/test_wrapper.py diff --git a/assemblyline_client/v4_client/wrapper.py b/assemblyline_client/v4_client/wrapper.py index 03a3166..19c66cd 100644 --- a/assemblyline_client/v4_client/wrapper.py +++ b/assemblyline_client/v4_client/wrapper.py @@ -8,7 +8,6 @@ from assemblyline_client.v4_client.module.submission import Submission from assemblyline_client.v4_client.module.workflow import Workflow - def wrapper_function(inner_func, args=[]): """\ Decorator for wrapper functions to provide docstrings for documentation. @@ -122,6 +121,7 @@ def get_extracted_files(self): query += f"{file['sha256']}" continue query += f"{file['sha256']} OR " + if query == "": return [] return self.search.file(query)['items'] except KeyError: @@ -292,7 +292,7 @@ def summary(self, *args, **kwargs): def tree(self, *args, **kwargs): return self.submission.tree(self['sid'], *args, **kwargs) - def get_files(self): + def get_submitted_files(self): """\ Get the list of files that were originally submitted @@ -330,7 +330,6 @@ def delete(self): def update(self, *args, **kwargs): return self.workflow.update(self['workflow_id'], *args, **kwargs) - WRAPPER_MAP = { 'file': FileWrapper, 'alert': AlertWrapper, diff --git a/test/test_wrapper.py b/test/test_wrapper.py new file mode 100644 index 0000000..771e0f1 --- /dev/null +++ b/test/test_wrapper.py @@ -0,0 +1,43 @@ +import random +from utils import random_id_from_collection + +def test_get_submissions(datastore, client): + subm = client.search.submission("*") + + sid = subm['items'][0]['sid'] + sha256 = subm['items'][0].full()['files'][0]['sha256'] + + file = client.search.file("sha256:{}".format(sha256))['items'][0] + + assert file.get_submissions()[0]['sid'] == sid + +def test_get_results(datastore, client): + query = client.search.result("*")['items'] + result = query[random.randint(0, len(query)-1)] + result_id = result['id'] + sha256 = result()['sha256'] + + file = client.search.file("sha256:{}".format(sha256))['items'][0] + results_wrapper = file.get_results() + + result_ids = [result['id'] for result in results_wrapper] + + assert result_id in result_ids + +def test_get_extracted_files(datastore, client): + query = client.search.result("*")['items'] + result = query[random.randint(0, len(query)-1)]() + + test_file = client.search.file(result['sha256'])['items'][0] + datastore_extracted_files = result['response']['extracted'] + + files_sha256 = [file['sha256'] for file in datastore_extracted_files] + extracted_sha256 = [file['sha256'] for file in test_file.get_extracted_files()] + + for sha256 in files_sha256: + assert sha256 in extracted_sha256 + +def test_get_submitted_files(datastore, client): + query = client.search.submission("*")['items'] + subm = query[random.randint(0, len(query)-1)] + assert len(subm.full()['files']) == len(subm.get_submitted_files())