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