def check(data): api = Diario('API - ID', 'API - SECRETS') result = api.search(data[1]) if result.error is None: return result.data elif result.json['error']['code'] == 406: api.upload(data[0]) result = {'prediction': 'U', 'status': 'Q', 'sha256': data[1]} return result else: result = {'prediction': 'E'} return result
class DiarioAnalyzer(Analyzer): def __init__(self): Analyzer.__init__(self) self.service = self.get_param('config.service', None, 'Service parameter is missing') self.client_id = self.get_param('config.client_id', None, 'Missing Client ID') self.secret = self.get_param('config.secret', None, 'Missing Secret') self.polling_interval = self.get_param('config.polling_interval', 60) self.api = Diario(self.client_id, self.secret) _predictions = { "M": "Malware", "G": "Goodware", "NM": "No Macros present", # Only applies to office documents "U": "Unknown" } _stages = { "A": "Analyzed", "Q": "Queued", "P": "Processing", "F": "Failed" } def check_response(self, document_hash): response = self.api.search(document_hash) if response.error: if response.error.code == 413: time.sleep(self.polling_interval) return self.check_response(document_hash) elif response.error.code in (406, 409): return dict(message=response.error.message) else: self.error(response.error) if response.data["status"] in ("P", "Q"): time.sleep(self.polling_interval) return self.check_response(document_hash) elif response.data["status"] == "F": self.error(response.data) data = response.data data["prediction"] = self._predictions.get(data["prediction"]) data["status"] = self._stages.get(data["status"]) return dict(data) def summary(self, raw): taxonomies = [] level = "info" namespace = "Diario" predicate = "GetReport" value = "Not Found" if self.service == "scan": predicate = "Scan" verdicts = { "Goodware": "safe", "Malware": "malicious", "Unknown": "suspicious", "No Macros present": "info", } if "sha256" in raw: value = raw["prediction"] level = verdicts.get(value) taxonomies.append( self.build_taxonomy(level, namespace, predicate, value)) return {"taxonomies": taxonomies} def run(self): if self.service == "scan": if self.data_type == "file": filepath = self.get_param('file', None, 'File is missing') response = self.api.upload(filepath) if response.error: self.error(response.error["message"]) data = response.data.get("hash", None) else: self.error("Data type has to be a File") return elif self.service == "get": # If we want to only get the report of a file we get the # SHA256 hash and check if there is a report if self.data_type == "file": filepath = self.get_param('file', None, 'File is missing') sha256_hash = hashlib.sha256() with open(filepath, "rb") as f: # Read and update hash string value in blocks of 4K for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) data = sha256_hash.hexdigest() elif self.data_type == "hash": data = self.get_param('data', None, 'Data is missing') else: self.error("Data type has to be a File or Hash") return else: self.error("Service doesn't exists") return self.report(self.check_response(data))