def perform_scan(self, poll_timeout):
'''
Wait for work from broker then perform the scan. If timeout occurs, no
scan is performed and no result is returned.
Arguments:
poll_timeout -- The amount of time to wait for work.
Returns:
The result of the scan or None if no scan was performed.
'''
from laikaboss.dispatch import Dispatch
from laikaboss.objectmodel import ScanResult, ExternalObject, ExternalVars
from laikaboss.util import log_result
# If task is found, perform scan
try:
logging.debug("Worker (%s): checking for work", self.identity)
tasks = dict(self.broker_poller.poll(poll_timeout))
if tasks.get(self.broker) == zmq.POLLIN:
logging.debug("Worker (%s): performing scan", self.identity)
# task should be in the following format
# ['', client_id, '', request_type, '', request]
# where:
# client_id -- ZMQ identifier of the client socket
# request_type -- The type of request (json/pickle/zlib)
# request -- Object to be scanned
task = self.broker.recv_multipart()
client_id = task[1]
if len(task) == 6:
request_type = task[3]
request = task[5]
if request_type in [REQ_TYPE_PICKLE, REQ_TYPE_PICKLE_ZLIB]:
#logging.debug("Worker: received work %s", str(task))
if request_type == REQ_TYPE_PICKLE_ZLIB:
externalObject = pickle.loads(zlib.decompress(request))
else:
externalObject = pickle.loads(request)
elif request_type in [REQ_TYPE_JSON, REQ_TYPE_JSON_ZLIB]:
if request_type == REQ_TYPE_JSON_ZLIB:
jsonRequest = json.loads(zlib.decompress(request))
else:
jsonRequest = json.loads(request)
# Set default values for our request just in case some were omitted
if not 'buffer' in jsonRequest:
jsonRequest['buffer'] = ''
else:
try:
jsonRequest['buffer'] = base64.b64decode(jsonRequest['buffer'])
except:
# This should never happen unless invalid input is given
jsonRequest['buffer'] = ''
if not 'filename' in jsonRequest:
jsonRequest['filename'] = ''
if not 'ephID' in jsonRequest:
jsonRequest['ephID'] = ''
if not 'uniqID' in jsonRequest:
jsonRequest['uniqID'] = ''
if not 'contentType' in jsonRequest:
jsonRequest['contentType'] = []
if not 'timestamp' in jsonRequest:
jsonRequest['timestamp'] = ''
if not 'source' in jsonRequest:
jsonRequest['source'] = ''
if not 'origRootUID' in jsonRequest:
jsonRequest['origRootUID'] = ''
if not 'extMetaData' in jsonRequest:
jsonRequest['extMetaData'] = {}
if not 'level' in jsonRequest:
jsonRequest['level'] = 2
externalVars = ExternalVars(filename=jsonRequest['filename'],
ephID=jsonRequest['ephID'],
uniqID=jsonRequest['uniqID'],
contentType=jsonRequest['contentType'],
timestamp=jsonRequest['timestamp'],
source=jsonRequest['source'],
origRootUID=jsonRequest['origRootUID'],
extMetaData=jsonRequest['extMetaData'])
externalObject = ExternalObject(buffer=jsonRequest['buffer'],
level=jsonRequest['level'],
externalVars=externalVars)
else:
return [client_id, '', 'INVALID REQUEST']
result = ScanResult(
source=externalObject.externalVars.source,
level=externalObject.level)
result.startTime = time.time()
try:
Dispatch(externalObject.buffer, result, 0,
externalVars=externalObject.externalVars)
except QuitScanException:
raise
except:
exc_type, exc_value, exc_traceback = sys.exc_info()
log_debug(
"exception on file: %s, detailed exception: %s" % (
externalObject.externalVars.filename,
repr(traceback.format_exception(
exc_type, exc_value, exc_traceback))))
if self.logresult:
log_result(result)
if request_type == REQ_TYPE_PICKLE_ZLIB:
result = zlib.compress(
pickle.dumps(result, pickle.HIGHEST_PROTOCOL))
elif request_type == REQ_TYPE_PICKLE:
result = pickle.dumps(result, pickle.HIGHEST_PROTOCOL)
elif request_type == REQ_TYPE_JSON_ZLIB:
result = zlib.compress(
json.dumps(result, cls=ResultEncoder))
elif request_type == REQ_TYPE_JSON:
result = json.dumps(result, cls=ResultEncoder)
return [client_id, '', result]
else:
return [client_id, '', 'INVALID REQUEST']
except zmq.ZMQError as zmqerror:
if "Interrupted system call" not in str(zmqerror):
logging.exception("Worker (%s): Received ZMQError", self.identity)
else:
logging.debug("Worker (%s): ZMQ interrupted by shutdown signal", self.identity)
return None
def run(self):
global CONFIG_PATH
config.init(path=CONFIG_PATH)
init_logging()
ret_value = 0
# Loop and accept messages from both channels, acting accordingly
while True:
next_task = self.task_queue.get()
if next_task is None:
# Poison pill means shutdown
self.task_queue.task_done()
logging.debug("%s Got poison pill" % (os.getpid()))
break
try:
with open(next_task) as nextfile:
file_buffer = nextfile.read()
except IOError:
logging.debug("Error opening: %s" % (next_task))
self.task_queue.task_done()
self.result_queue.put(answer)
continue
resultJSON = ""
try:
# perform the work
result = ScanResult()
result.source = SOURCE
result.startTime = time.time()
result.level = level_metadata
myexternalVars = ExternalVars(filename=next_task,
source=SOURCE,
ephID=EPHID,
extMetaData=EXT_METADATA)
Dispatch(file_buffer, result, 0, externalVars=myexternalVars, extScanModules=SCAN_MODULES)
resultJSON = getJSON(result)
if SAVE_PATH:
rootObject = getRootObject(result)
UID_SAVE_PATH = "%s/%s" % (SAVE_PATH, get_scanObjectUID(rootObject))
if not os.path.exists(UID_SAVE_PATH):
try:
os.makedirs(UID_SAVE_PATH)
except (OSError, IOError) as e:
error("\nERROR: unable to write to %s...\n" % (UID_SAVE_PATH))
raise
for uid, scanObject in result.files.iteritems():
with open("%s/%s" % (UID_SAVE_PATH, uid), "wb") as f:
f.write(scanObject.buffer)
if scanObject.filename and scanObject.depth != 0:
linkPath = "%s/%s" % (UID_SAVE_PATH, scanObject.filename.replace("/","_"))
if not os.path.lexists(linkPath):
os.symlink("%s" % (uid), linkPath)
elif scanObject.filename:
filenameParts = scanObject.filename.split("/")
os.symlink("%s" % (uid), "%s/%s" % (UID_SAVE_PATH, filenameParts[-1]))
with open("%s/%s" % (UID_SAVE_PATH, "result.json"), "wb") as f:
f.write(resultJSON)
if LOG_RESULT:
log_result(result)
if LOG_JSON:
LOCAL_PATH = LOG_JSON
with open(LOCAL_PATH, "ab") as f:
f.write(resultJSON + "\n")
except:
logging.exception("Scan worker died, shutting down")
ret_value = 1
break
finally:
self.task_queue.task_done()
self.result_queue.put(zlib.compress(resultJSON))
close_modules()
return ret_value
def perform_scan(self, poll_timeout):
'''
Wait for work from broker then perform the scan. If timeout occurs, no
scan is performed and no result is returned.
Arguments:
poll_timeout -- The amount of time to wait for work.
Returns:
The result of the scan or None if no scan was performed.
'''
from laikaboss.dispatch import Dispatch
from laikaboss.objectmodel import ScanResult, ExternalObject, ExternalVars
from laikaboss.util import log_result
# If task is found, perform scan
try:
logging.debug("Worker (%s): checking for work", self.identity)
tasks = dict(self.broker_poller.poll(poll_timeout))
if tasks.get(self.broker) == zmq.POLLIN:
logging.debug("Worker (%s): performing scan", self.identity)
# task should be in the following format
# ['', client_id, '', request_type, '', request]
# where:
# client_id -- ZMQ identifier of the client socket
# request_type -- The type of request (json/pickle/zlib)
# request -- Object to be scanned
task = self.broker.recv_multipart()
client_id = task[1]
if len(task) == 6:
request_type = task[3]
request = task[5]
if request_type in [REQ_TYPE_PICKLE, REQ_TYPE_PICKLE_ZLIB]:
#logging.debug("Worker: received work %s", str(task))
if request_type == REQ_TYPE_PICKLE_ZLIB:
externalObject = pickle.loads(zlib.decompress(request))
else:
externalObject = pickle.loads(request)
elif request_type in [REQ_TYPE_JSON, REQ_TYPE_JSON_ZLIB]:
if request_type == REQ_TYPE_JSON_ZLIB:
jsonRequest = json.loads(zlib.decompress(request))
else:
jsonRequest = json.loads(request)
# Set default values for our request just in case some were omitted
if not 'buffer' in jsonRequest:
jsonRequest['buffer'] = ''
else:
try:
jsonRequest['buffer'] = base64.b64decode(jsonRequest['buffer'])
except:
# This should never happen unless invalid input is given
jsonRequest['buffer'] = ''
if not 'filename' in jsonRequest:
jsonRequest['filename'] = ''
if not 'ephID' in jsonRequest:
jsonRequest['ephID'] = ''
if not 'uniqID' in jsonRequest:
jsonRequest['uniqID'] = ''
if not 'contentType' in jsonRequest:
jsonRequest['contentType'] = []
if not 'timestamp' in jsonRequest:
jsonRequest['timestamp'] = ''
if not 'source' in jsonRequest:
jsonRequest['source'] = ''
if not 'origRootUID' in jsonRequest:
jsonRequest['origRootUID'] = ''
if not 'extMetaData' in jsonRequest:
jsonRequest['extMetaData'] = {}
if not 'level' in jsonRequest:
jsonRequest['level'] = 2
externalVars = ExternalVars(filename=jsonRequest['filename'],
ephID=jsonRequest['ephID'],
uniqID=jsonRequest['uniqID'],
contentType=jsonRequest['contentType'],
timestamp=jsonRequest['timestamp'],
source=jsonRequest['source'],
origRootUID=jsonRequest['origRootUID'],
extMetaData=jsonRequest['extMetaData'])
externalObject = ExternalObject(buffer=jsonRequest['buffer'],
level=jsonRequest['level'],
externalVars=externalVars)
else:
return [client_id, '', 'INVALID REQUEST']
result = ScanResult(
source=externalObject.externalVars.source,
level=externalObject.level)
result.startTime = time.time()
try:
Dispatch(externalObject.buffer, result, 0,
externalVars=externalObject.externalVars)
except QuitScanException:
raise
except:
exc_type, exc_value, exc_traceback = sys.exc_info()
log_debug(
"exception on file: %s, detailed exception: %s" % (
externalObject.externalVars.filename,
repr(traceback.format_exception(
exc_type, exc_value, exc_traceback))))
if self.logresult:
log_result(result)
if request_type == REQ_TYPE_PICKLE_ZLIB:
result = zlib.compress(
pickle.dumps(result, pickle.HIGHEST_PROTOCOL))
elif request_type == REQ_TYPE_PICKLE:
result = pickle.dumps(result, pickle.HIGHEST_PROTOCOL)
elif request_type == REQ_TYPE_JSON_ZLIB:
result = zlib.compress(
json.dumps(result, cls=ResultEncoder))
elif request_type == REQ_TYPE_JSON:
result = json.dumps(result, cls=ResultEncoder)
return [client_id, '', result]
else:
return [client_id, '', 'INVALID REQUEST']
except zmq.ZMQError as zmqerror:
if "Interrupted system call" not in str(zmqerror):
logging.exception("Worker (%s): Received ZMQError", self.identity)
else:
logging.debug("Worker (%s): ZMQ interrupted by shutdown signal", self.identity)
return None
def run(self):
global CONFIG_PATH
config.init(path=CONFIG_PATH)
init_logging()
ret_value = 0
# Loop and accept messages from both channels, acting accordingly
while True:
next_task = self.task_queue.get()
if next_task is None:
# Poison pill means shutdown
self.task_queue.task_done()
logging.debug("%s Got poison pill" % (os.getpid()))
break
try:
with open(next_task) as nextfile:
file_buffer = nextfile.read()
except IOError:
logging.debug("Error opening: %s" % (next_task))
self.task_queue.task_done()
self.result_queue.put(answer)
continue
resultJSON = ""
try:
# perform the work
result = ScanResult()
result.source = SOURCE
result.startTime = time.time()
result.level = level_metadata
myexternalVars = ExternalVars(filename=next_task,
source=SOURCE,
ephID=EPHID,
extMetaData=EXT_METADATA)
Dispatch(file_buffer, result, 0, externalVars=myexternalVars, extScanModules=SCAN_MODULES)
resultJSON = getJSON(result)
if SAVE_PATH:
rootObject = getRootObject(result)
UID_SAVE_PATH = "%s/%s" % (SAVE_PATH, get_scanObjectUID(rootObject))
if not os.path.exists(UID_SAVE_PATH):
try:
os.makedirs(UID_SAVE_PATH)
except (OSError, IOError) as e:
error("\nERROR: unable to write to %s...\n" % (UID_SAVE_PATH))
raise
for uid, scanObject in result.files.iteritems():
with open("%s/%s" % (UID_SAVE_PATH, uid), "wb") as f:
f.write(scanObject.buffer)
if scanObject.filename and scanObject.depth != 0:
linkPath = "%s/%s" % (UID_SAVE_PATH, scanObject.filename.replace("/","_"))
if not os.path.lexists(linkPath):
os.symlink("%s" % (uid), linkPath)
elif scanObject.filename:
filenameParts = scanObject.filename.split("/")
os.symlink("%s" % (uid), "%s/%s" % (UID_SAVE_PATH, filenameParts[-1]))
with open("%s/%s" % (UID_SAVE_PATH, "result.json"), "wb") as f:
f.write(resultJSON)
if LOG_RESULT:
log_result(result)
except:
logging.exception("Scan worker died, shutting down")
ret_value = 1
break
finally:
self.task_queue.task_done()
self.result_queue.put(zlib.compress(resultJSON))
close_modules()
return ret_value
