Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Al-705 Assemblyline Client Wrapper #56

Open
wants to merge 25 commits into
base: master
Choose a base branch
from
276 changes: 220 additions & 56 deletions assemblyline_client/v4_client/module/search/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +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, BaseWrapper


class Search(object):
Expand All @@ -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)

Expand All @@ -37,10 +44,23 @@ 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))

def alert(self, query, filters=None, fl=None, offset=0, rows=25, sort=None, timeout=None,
use_archive=False, track_total_hits=None):
data = self._connection.post(path, data=json.dumps(kwargs))
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):
"""\
Search alerts with a lucene query.

Expand All @@ -59,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.

Expand All @@ -83,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.

Expand All @@ -107,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.

Expand All @@ -131,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.

Expand All @@ -155,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.

Expand All @@ -179,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.

Expand All @@ -203,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.

Expand All @@ -227,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.

Expand All @@ -251,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)
Loading