def _do_stream(self, bucket, query, **kwargs):
if bucket not in SEARCHABLE:
raise ClientError("Bucket %s is not searchable" % bucket, 400)
for arg in list(kwargs.keys()):
if arg in INVALID_STREAM_SEARCH_PARAMS:
raise ClientError(
"The following parameters cannot be used with stream search: %s",
", ".join(INVALID_STREAM_SEARCH_PARAMS))
kwargs.update({'rows': str(self._page_size), 'deep_paging_id': '*'})
yield_done = False
items = []
lock = threading.Lock()
sf_t = threading.Thread(target=self._auto_fill,
args=[items, lock, bucket, query],
kwargs=kwargs)
sf_t.setDaemon(True)
sf_t.start()
while not yield_done:
try:
with lock:
item = items.pop(0)
yield item
except IndexError:
if not sf_t.is_alive() and len(items) == 0:
yield_done = True
time.sleep(0.01)
def _do_stream(self, bucket, query, *args, **kwargs):
if bucket not in SEARCHABLE:
raise ClientError("Bucket %s is not searchable" % bucket, 0)
for arg in list(args) + list(kwargs.items()):
if arg[0] in INVALID_STREAM_SEARCH_PARAMS:
raise ClientError(
"The following parameters cannot be used with stream search: %s",
", ".join(INVALID_STREAM_SEARCH_PARAMS))
args = list(args) + [('sort', '_yz_id asc'),
('rows', str(self._page_size)),
('cursorMark', '*')]
yield_done = False
items = []
lock = threading.Lock()
sf_t = threading.Thread(target=self._auto_fill,
args=[items, lock, bucket, query] + args,
kwargs=kwargs)
sf_t.setDaemon(True)
sf_t.start()
while not yield_done:
try:
with lock:
item = items.pop(0)
yield item
except IndexError:
if not sf_t.is_alive() and len(items) == 0:
yield_done = True
time.sleep(0.01)
def __call__(self,
path=None,
contents=None,
url=None,
sha256=None,
fname=None,
params=None,
srv_spec=None):
"""\
Submit a file to be dispatched.
Required (one of)
contents: Content of the file to scan
path : Path/name of file. (string)
sha256 : Sha256 of the file to scan
url : Url to scan
Optional
fname : Name of the file to scan
params : Additional submission parameters. (dict)
srv_spec: Service specific submission parameters. (dict)
If contents are provided, the path is used as metadata only.
"""
if contents is None and path:
if os.path.exists(path):
with open(path, 'rb') as f:
contents = f.read()
else:
raise ClientError('File does not exist "%s"' % path, 400)
if contents:
request = {
'binary': b64encode(contents).decode('ascii'),
'name': fname or os.path.basename(path),
}
elif url:
request = {
'url': url,
'name': fname or os.path.basename(url).split("?")[0],
}
elif sha256:
request = {
'sha256': sha256,
'name': fname or sha256,
}
else:
raise ClientError(
'You need to provide at least content, a path, a url or a sha256',
400)
if params:
request['params'] = params
if srv_spec:
request['srv_spec'] = srv_spec
return self._connection.post(_path('submit'), data=dumps(request))
def request(self, func, path, process, **kw):
self.debug(path)
# Apply default timeout parameter if not passed elsewhere
kw.setdefault('timeout', self.default_timeout)
retries = 0
with warnings.catch_warnings():
if self.silence_warnings:
warnings.simplefilter('ignore')
while self.max_retries < 1 or retries <= self.max_retries:
if retries:
time.sleep(min(2, 2**(retries - 7)))
response = func('/'.join((self.server, path)), **kw)
if 'XSRF-TOKEN' in response.cookies:
self.session.headers.update(
{'X-XSRF-TOKEN': response.cookies['XSRF-TOKEN']})
if response.ok:
return process(response)
elif response.status_code == 401:
try:
resp_data = response.json()
if resp_data[
"api_error_message"] == "Authentication required":
self._authenticate()
else:
raise ClientError(
resp_data["api_error_message"],
response.status_code,
api_version=resp_data["api_server_version"],
api_response=resp_data["api_response"])
except Exception as e:
if isinstance(e, ClientError):
raise
raise ClientError(response.content,
response.status_code)
elif response.status_code not in (502, 503, 504):
try:
resp_data = response.json()
raise ClientError(
resp_data["api_error_message"],
response.status_code,
api_version=resp_data["api_server_version"],
api_response=resp_data["api_response"])
except Exception as e:
if isinstance(e, ClientError):
raise
raise ClientError(response.content,
response.status_code)
retries += 1
def _error_callback(self, data):
try:
self._sio.disconnect()
except Exception:
pass
raise ClientError(data['err_msg'], data['status_code'])
def listen_on_submissions(self,
completed_callback=None,
ingested_callback=None,
received_callback=None,
started_callback=None,
timeout=None):
"""\
Listen to the various submission messages in the system and call the callback for each of them
Required:
completed_callback : Callback function for when submission completed messages are received
ingested_callback : Callback function for when submission ingested messages are received
received_callback : Callback function for when submission received messages are received
started_callback : Callback function for when submission started messages are received
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if ingested_callback is None and received_callback is None and \
completed_callback is None and started_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketio.Client(ssl_verify=False)
self._stop_on_warning.set_sio(self._sio)
if ingested_callback is not None:
self._sio.on("SubmissionIngested",
ingested_callback,
namespace='/submissions')
if received_callback is not None:
self._sio.on("SubmissionReceived",
received_callback,
namespace='/submissions')
if completed_callback is not None:
self._sio.on("SubmissionCompleted",
completed_callback,
namespace='/submissions')
if started_callback is not None:
self._sio.on("SubmissionStarted",
started_callback,
namespace='/submissions')
self._sio.connect(self._server,
namespaces=['/submissions'],
headers=deepcopy(self._header))
self._sio.emit('monitor', {
"status": "start",
"client": "assemblyline_client"
},
namespace='/submissions')
if timeout is None:
self._sio.wait()
else:
self._sio.sleep(timeout)
self._sio.disconnect()
def _do_search(self, bucket, query, *args, **kwargs):
if bucket not in SEARCHABLE:
raise ClientError("Bucket %s is not searchable" % bucket, 0)
args = [('df', 'text')] + list(args) + list(kwargs.items())
params = _join_params(query, args)
path = '?q='.join((_path('search/advanced', bucket), params))
return self._connection.get(path)
def listen_on_watch_queue(self,
wq,
result_callback=None,
error_callback=None,
timeout=None):
"""\
Listen to the various messages of a currently running submission's watch queue
Required:
wq : ID of the watch queue to listen for
result_callback : Callback function when receiveing a result cache key
error_callback : Callback function when receiveing a error cache key
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if result_callback is None and error_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketio.Client(ssl_verify=False)
self._stop_on_warning.set_sio(self._sio)
if result_callback:
self._sio.on("cachekey",
result_callback,
namespace='/live_submission')
if error_callback:
self._sio.on("cachekeyerr",
error_callback,
namespace='/live_submission')
self._sio.on("stop", self._stop_callback, namespace='/live_submission')
self._sio.on("error",
self._error_callback,
namespace='/live_submission')
self._sio.connect(self._server,
namespaces=['/live_submission'],
headers=deepcopy(self._header))
self._sio.emit('listen', {
"status": "start",
"client": "assemblyline_client",
"wq_id": wq,
'from_start': True
},
namespace="/live_submission")
if timeout is None:
self._sio.wait()
else:
self._sio.sleep(timeout)
self._sio.disconnect()
def __init__( # pylint: disable=R0913
self, server, auth, cert, debug, headers, retries,
silence_warnings, apikey, verify, timeout):
self.auth = auth
self.apikey = apikey
self.debug = debug
self.is_v4 = False
self.max_retries = retries
self.server = server
self.silence_warnings = silence_warnings
self.verify = verify
self.default_timeout = timeout
session = requests.Session()
session.headers.update({'content-type': 'application/json'})
session.verify = verify
if cert:
session.cert = cert
if headers:
session.headers.update(headers)
self.session = session
try:
auth_session_detail = self._authenticate()
except requests.exceptions.SSLError as ssle:
raise ClientError(
"Client could not connect to the server "
"due to the following SSLError: %s" % ssle, 495)
session.timeout = auth_session_detail['session_duration']
r = self.request(self.session.get, 'api/', convert_api_output)
if not isinstance(r, list) or not set(r).intersection(SUPPORTED_APIS):
raise ClientError(
"Supported APIS (%s) are not available" % SUPPORTED_APIS, 400)
def _do_grouped(self, bucket, field, **kwargs):
if bucket not in SEARCHABLE:
raise ClientError("Bucket %s is not searchable" % bucket, 400)
filters = kwargs.pop('filters', None)
if filters is not None:
if isinstance(filters, str):
filters = [filters]
filters = [('filters', fq) for fq in filters]
kwargs = {k: v for k, v in kwargs.items() if v is not None and k != 'filters'}
if filters is not None:
kwargs['params_tuples'] = filters
path = api_path('search', 'grouped', bucket, field, **kwargs)
return self._connection.get(path)
def listen_on_alerts_messages(self,
alert_created_callback=None,
alert_updated_callback=None,
timeout=None):
"""\
Listen to the various alerts created messages in the system and call the callback for each alerts
Required:
alert_created_callback : Callback function for when alerts created messages are received
alert_updated_callback : Callback function for when alerts updated messages are received
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if alert_created_callback is None and alert_updated_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketio.Client(ssl_verify=False)
self._stop_on_warning.set_sio(self._sio)
if alert_created_callback:
self._sio.on("AlertCreated",
alert_created_callback,
namespace='/alerts')
if alert_updated_callback:
self._sio.on("AlertUpdated",
alert_updated_callback,
namespace='/alerts')
self._sio.connect(self._server,
namespaces=['/alerts'],
headers=deepcopy(self._header))
self._sio.emit('alert', {
"status": "start",
"client": "assemblyline_client"
},
namespace='/alerts')
if timeout is None:
self._sio.wait()
else:
self._sio.sleep(timeout)
self._sio.disconnect()
def listen_on_watch_queue(self,
wq,
result_callback=None,
error_callback=None):
"""\
Listen to the various messages of a currently running submission's watch queue
Required:
wq : ID of the watch queue to listen for
result_callback : Callback function when receiveing a result cache key
error_callback : Callback function when receiveing a error cache key
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if result_callback is None and error_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketIO_client.SocketIO(self._server,
port=self._port,
headers=self._header,
verify=self._verify)
self._stop_on_warning.set_sio(self._sio)
if result_callback:
self._sio.on("cachekey", result_callback)
if error_callback:
self._sio.on("cachekeyerr", error_callback)
self._sio.on("stop", self._stop_callback)
self._sio.on("error", self._error_callback)
self._sio.emit(
'listen', {
"status": "start",
"client": "assemblyline_client",
"wq_id": wq,
'from_start': True
})
self._sio.wait()
def listen_on_dashboard_messages(self,
dispatcher_msg_callback=None,
ingest_msg_callback=None,
service_msg_callback=None):
"""\
Listen to the various messages you would find on the UI dashboard.
Required (one of):
dispatcher_msg_callback : Callback function when a dispatcher message is received
ingest_msg_callback : Callback function when a ingest message is received
service_msg_callback : Callback function when a service message is received
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if dispatcher_msg_callback is None and ingest_msg_callback is None and service_msg_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketIO_client.SocketIO(self._server,
port=self._port,
headers=self._header,
verify=self._verify)
self._stop_on_warning.set_sio(self._sio)
if dispatcher_msg_callback:
self._sio.on("DispHeartbeat", dispatcher_msg_callback)
if ingest_msg_callback:
self._sio.on("IngestHeartbeat", ingest_msg_callback)
if service_msg_callback:
self._sio.on("SvcHeartbeat", service_msg_callback)
self._sio.emit('monitor', {
"status": "start",
"client": "assemblyline_client"
})
self._sio.wait()
def _do_fields(self, bucket):
if bucket not in SEARCHABLE:
raise ClientError("Bucket %s is not searchable" % bucket, 400)
path = api_path('search', 'fields', bucket)
return self._connection.get(path)
def __call__(self,
path=None,
content=None,
url=None,
sha256=None,
fname=None,
params=None,
metadata=None,
alert=False,
nq=None,
nt=None,
ingest_type='AL_CLIENT'):
"""\
Submit a file to the ingestion queue.
Required (one of)
content : Content of the file to scan (byte array)
path : Path/name of file (string)
sha256 : Sha256 of the file to scan (string)
url : Url to scan (string)
Optional
alert : Create an alert if score above alert threshold. (boolean)
fname : Name of the file to scan (string)
metadata : Metadata to include with submission. (dict)
nq : Notification queue name. (string)
nt : Notification threshold. (int)
params : Additional submission parameters. (dict)
ingest_type: Ingestion type, one word to describe how the data is ingested. Default: AL_CLIENT (string)
If content is provided, the path is used as metadata only.
"""
if content:
fd, path = tempfile.mkstemp()
with os.fdopen(fd, 'wb') as fh:
if isinstance(content, str):
content = content.encode()
fh.write(content)
files = {}
if path:
if os.path.exists(path):
files = {'bin': open(path, 'rb')}
else:
raise ClientError('File does not exist "%s"' % path, 400)
request = {'name': fname or os.path.basename(path)}
elif url:
request = {
'url': url,
'name': fname or os.path.basename(url).split("?")[0],
}
elif sha256:
request = {
'sha256': sha256,
'name': fname or sha256,
}
else:
raise ClientError(
'You need to provide at least content, a path, a url or a sha256',
400)
request.update({
'metadata': {},
'type': ingest_type,
})
if alert:
request['generate_alert'] = bool(alert)
if metadata:
request['metadata'].update(metadata)
if nq:
request['notification_queue'] = nq
if nt:
request['notification_threshold'] = int(nt)
if params:
request['params'] = params
if files:
data = {'json': dumps(request)}
headers = {'content-type': None}
else:
data = dumps(request)
headers = None
return self._connection.post(api_path('ingest'),
data=data,
files=files,
headers=headers)
def listen_on_status_messages(self,
alerter_msg_callback=None,
archive_msg_callback=None,
dispatcher_msg_callback=None,
expiry_msg_callback=None,
ingest_msg_callback=None,
scaler_msg_callback=None,
scaler_status_msg_callback=None,
service_msg_callback=None,
timeout=None):
"""\
Listen to the various status messages you would find on the UI dashboard.
Required (one of):
alerter_msg_callback : Callback function when an alerter message is received
archive_msg_callback : Callback function when an archive message is received
dispatcher_msg_callback : Callback function when a dispatcher message is received
expiry_msg_callback : Callback function when an expiry message is received
ingest_msg_callback : Callback function when an ingest message is received
scaler_msg_callback : Callback function when an scaler message is received
scaler_status_msg_callback : Callback function when a scaler status message is received
service_msg_callback : Callback function when a service message is received
This function wait indefinitely and calls the appropriate callback for each messages returned
"""
if dispatcher_msg_callback is None and ingest_msg_callback is None and service_msg_callback is None and \
alerter_msg_callback is None and expiry_msg_callback is None and scaler_msg_callback is None and \
archive_msg_callback is None and scaler_status_msg_callback is None:
raise ClientError(
"At least one of the callbacks needs to be defined...", 400)
self._sio = socketio.Client(ssl_verify=False)
self._stop_on_warning.set_sio(self._sio)
if alerter_msg_callback:
self._sio.on("AlerterHeartbeat",
alerter_msg_callback,
namespace='/status')
if archive_msg_callback:
self._sio.on("ArchiveHeartbeat",
archive_msg_callback,
namespace='/status')
if dispatcher_msg_callback:
self._sio.on("DispatcherHeartbeat",
dispatcher_msg_callback,
namespace='/status')
if expiry_msg_callback:
self._sio.on("ExpiryHeartbeat",
expiry_msg_callback,
namespace='/status')
if ingest_msg_callback:
self._sio.on("IngestHeartbeat",
ingest_msg_callback,
namespace='/status')
if scaler_msg_callback:
self._sio.on("ScalerHeartbeat",
scaler_msg_callback,
namespace='/status')
if scaler_status_msg_callback:
self._sio.on("ScalerStatusHeartbeat",
scaler_status_msg_callback,
namespace='/status')
if service_msg_callback:
self._sio.on("ServiceHeartbeat",
service_msg_callback,
namespace='/status')
self._sio.connect(self._server,
namespaces=['/status'],
headers=deepcopy(self._header))
self._sio.emit('monitor', {
"status": "start",
"client": "assemblyline_client"
},
namespace='/status')
if timeout is None:
self._sio.wait()
else:
self._sio.sleep(timeout)
self._sio.disconnect()
def __call__(self,
path=None,
content=None,
url=None,
sha256=None,
fname=None,
params=None,
metadata=None):
"""\
Submit a file to be dispatched.
Required (one of)
content : Content of the file to scan
path : Path/name of file. (string)
sha256 : Sha256 of the file to scan
url : Url to scan
Optional
fname : Name of the file to scan
metadata : Metadata to include with submission. (dict)
params : Additional submission parameters. (dict)
If content is provided, the path is used as metadata only.
"""
if content:
fd, path = tempfile.mkstemp()
with os.fdopen(fd, 'wb') as fh:
if isinstance(content, str):
content = content.encode()
fh.write(content)
files = {}
if path:
if os.path.exists(path):
files = {'bin': open(path, 'rb')}
else:
raise ClientError('File does not exist "%s"' % path, 400)
request = {'name': fname or os.path.basename(path)}
elif url:
request = {
'url': url,
'name': fname or os.path.basename(url).split("?")[0],
}
elif sha256:
request = {
'sha256': sha256,
'name': fname or sha256,
}
else:
raise ClientError(
'You need to provide at least content, a path, a url or a sha256',
400)
if params:
request['params'] = params
if metadata:
request['metadata'] = metadata
if files:
data = {'json': dumps(request)}
headers = {'content-type': None}
else:
data = dumps(request)
headers = None
return self._connection.post(api_path('submit'),
data=data,
files=files,
headers=headers)
def __call__(
self,
path=None,
alert=False,
contents=None,
metadata=None,
nq=None,
nt=None,
params=None,
srv_spec=None,
fname=None,
url=None,
sha256=None,
ingest_type='AL_CLIENT' # pylint: disable=W0622
):
"""\
Submit a file to the ingestion queue.
Required (one of)
contents: Content of the file to scan
path : Path/name of file. (string)
sha256 : Sha256 of the file to scan
url : Url to scan
Optional
alert : Create an alert if score above alert threshold. (boolean)
contents: File contents. (string)
fname : Name of the file to scan
metadata: Metadata to include with submission. (dict)
nq : Notification queue name. (string)
nt : Notification threshold. (int)
params : Additional submission parameters. (dict)
srv_spec: Service-specific parameters. (dict)
If contents are provided, the path is used as metadata only.
"""
if contents is None and path:
if os.path.exists(path):
with open(path, 'rb') as f:
contents = f.read()
else:
raise ClientError('File does not exist "%s"' % path, 400)
if contents:
request = {
'binary': b64encode(contents).decode('ascii'),
'name': fname or os.path.basename(path),
}
elif url:
request = {
'url': url,
'name': fname or os.path.basename(url).split("?")[0],
}
elif sha256:
request = {
'sha256': sha256,
'name': fname or sha256,
}
else:
raise ClientError(
'You need to provide at least content, a path, a url or a sha256',
400)
request.update({
'metadata': {
'filename': request['name']
},
'type': ingest_type,
})
if alert:
request['generate_alert'] = bool(alert)
if metadata:
request['metadata'].update(metadata)
if nq:
request['notification_queue'] = nq
if nt:
request['notification_threshold'] = int(nt)
if params:
request['params'] = params
if srv_spec:
request['srv_spec'] = srv_spec
return self._connection.post(_path('ingest'), data=dumps(request))