def test013_load_from_ext_id_multiple(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().options().filter
m_filter.side_effect = MultipleResultsFound()
with self.assertRaises(IrmaDatabaseError):
Scan.load_from_ext_id(external_id, m_session)
def test013_load_from_ext_id_multiple(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().filter
m_filter.side_effect = MultipleResultsFound()
with self.assertRaises(IrmaDatabaseError):
Scan.load_from_ext_id(external_id, m_session)
def test012_load_from_ext_id_none(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().filter
m_filter.side_effect = NoResultFound()
with self.assertRaises(IrmaDatabaseResultNotFound):
Scan.load_from_ext_id(external_id, m_session)
def test012_load_from_ext_id_none(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().options().filter
m_filter.side_effect = NoResultFound()
with self.assertRaises(IrmaDatabaseResultNotFound):
Scan.load_from_ext_id(external_id, m_session)
def test011_load_from_ext_id(self):
m_session = MagicMock()
m_query = MagicMock()
m_session.query.return_value = m_query
external_id = "whatever"
Scan.load_from_ext_id(external_id, m_session)
m_filter = m_query.filter
m_filter.assert_called_once()
def test011_load_from_ext_id(self):
m_session = MagicMock()
m_query = MagicMock()
m_session.query.return_value = m_query
external_id = "whatever"
Scan.load_from_ext_id(external_id, m_session)
m_filter = m_query.options().filter
m_filter.assert_called_once()
def add_files(scanid, db):
""" Attach a file to a scan.
The request should be performed using a POST request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
if len(request.files) == 0:
raise ValueError("No files uploaded")
files = {}
for f in request.files:
upfile = request.files.get(f)
filename = decode_utf8(upfile.raw_filename)
data = upfile.file
files[filename] = data
scan_ctrl.add_files(scan, files, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def list(db):
""" Get a list of all scans which have been launched.
"""
try:
# Get values from query or set to default
offset = request.query.get("offset", default=0)
offset = int(offset)
limit = request.query.get("limit", default=5)
limit = int(limit)
log.debug("offset %d limit %d", offset, limit)
base_query = Scan.query_joined(db)
items = base_query.limit(limit).offset(offset).all()
if offset == 0 and len(items) < limit:
total = len(items)
else:
total = base_query.count()
log.debug("found %d scans", total)
response.content_type = "application/json; charset=UTF-8"
return {
"total": total,
"offset": offset,
"limit": limit,
"data": scan_schema.dump(items, many=True).data,
}
except Exception as e:
log.exception(e)
process_error(e)
def set_result(scanid, file_hash, probe, result):
with session_transaction() as session:
scan = Scan.load_from_ext_id(scanid, session=session)
fws = scan.get_filewebs_by_sha256(file_hash)
if len(fws) == 0:
log.error("file %s not found in scan", file_hash)
return
fws_file = File.load_from_sha256(file_hash, session)
fws_file.timestamp_last_scan = compat.timestamp()
fws_file.update(['timestamp_last_scan'], session=session)
sanitized_res = _sanitize_res(result)
# update results for all files with same sha256
for fw in fws:
# Update main reference results with fresh results
pr = _fetch_probe_result(fw, probe)
_update_ref_results(fw, fw.file, pr)
fw.file.update(session=session)
# fill ProbeResult with probe raw results
pr.doc = sanitized_res
pr.status = sanitized_res.get('status', None)
s_type = sanitized_res.get('type', None)
pr.type = IrmaProbeType.normalize(s_type)
pr.update(session=session)
probedone = []
for fw_pr in fw.probe_results:
if fw_pr.doc is not None:
probedone.append(fw_pr.name)
log.info("scanid: %s result from %s probedone %s", scanid, probe,
probedone)
is_finished(scanid)
def launch_asynchronous(scanid):
log.debug("scanid: %s", scanid)
with session_transaction() as session:
scan = Scan.load_from_ext_id(scanid, session=session)
IrmaScanStatus.filter_status(scan.status, IrmaScanStatus.ready,
IrmaScanStatus.ready)
scan_request = _create_scan_request(scan.files_web,
scan.get_probelist(),
scan.mimetype_filtering)
scan_request = _add_empty_results(scan.files_web, scan_request, scan,
session)
# Nothing to do
if scan_request.nb_files == 0:
scan.set_status(IrmaScanStatus.finished)
session.commit()
log.warning("scanid: %s finished nothing to do", scanid)
return
try:
upload_list = list()
for file in scan.files:
upload_list.append(file.path)
ftp_ctrl.upload_scan(scanid, upload_list)
except IrmaFtpError as e:
log.error("scanid: %s ftp upload error %s", scanid, str(e))
scan.set_status(IrmaScanStatus.error_ftp_upload)
session.commit()
return
# launch new celery scan task on brain
celery_brain.scan_launch(scanid, scan_request.to_dict())
scan.set_status(IrmaScanStatus.uploaded)
session.commit()
log.info("scanid: %s uploaded", scanid)
return
def add_files(scanid, db):
""" Attach a file to a scan.
The request should be performed using a POST request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
if len(request.files) == 0:
raise ValueError("No files uploaded")
files = {}
for f in request.files:
upfile = request.files.get(f)
filename = decode_utf8(upfile.raw_filename)
data = upfile.file
files[filename] = data
scan_ctrl.add_files(scan, files, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def launch(scanid, db):
""" Launch a scan.
The request should be performed using a POST request method.
"""
try:
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
probes = None
# handle scan parameter / cached results: "force"
if 'force' in request.json and request.json.get('force'):
scan.force = True
db.commit()
# V1 retro compatibility
scan.mimetype_filtering = False
scan.file_resubmit = False
# handle scan parameter / probelist: "probes"
if 'probes' in request.json:
probes = request.json.get('probes').split(',')
msg = "scanid: %s Force %s MimeF %s"
msg += "Resub %s Probes %s"
log.debug(msg, scanid, scan.force, scan.mimetype_filtering,
scan.resubmit_files, probes)
scan_ctrl.check_probe(scan, probes, db)
# launch_asynchronous scan via frontend task
celery_frontend.scan_launch(scanid)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def launch(scanid, db):
""" Launch a scan.
The request should be performed using a POST request method.
"""
try:
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
probes = None
# handle scan parameter / cached results: "force"
if 'force' in request.json and request.json.get('force'):
scan.force = True
db.commit()
# V1 retro compatibility
scan.mimetype_filtering = False
scan.file_resubmit = False
# handle scan parameter / probelist: "probes"
if 'probes' in request.json:
probes = request.json.get('probes').split(',')
msg = "scanid: %s Force %s MimeF %s"
msg += "Resub %s Probes %s"
log.debug(msg, scanid, scan.force, scan.mimetype_filtering,
scan.resubmit_files, probes)
scan_ctrl.check_probe(scan, probes, db)
# launch_asynchronous scan via frontend task
celery_frontend.scan_launch(scanid)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def new(db):
""" Create a new scan.
The request should be performed using a POST request method.
"""
try:
ip = request.remote_addr
scan = Scan(compat.timestamp(), ip)
db.add(scan)
scan.set_status(IrmaScanStatus.empty)
db.commit()
log.debug("scanid: %s", scan.external_id)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def new(db):
""" Create a new scan.
The request should be performed using a POST request method.
"""
try:
ip = request.remote_addr
scan = Scan(compat.timestamp(), ip)
db.add(scan)
scan.set_status(IrmaScanStatus.empty)
db.commit()
log.debug("scanid: %s", scan.external_id)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def is_finished(scanid):
with session_transaction() as session:
scan = Scan.load_from_ext_id(scanid, session=session)
if scan.finished() and scan.status != IrmaScanStatus.finished:
scan.set_status(IrmaScanStatus.finished)
session.commit()
# launch flush celery task on brain
log.debug("scanid: %s calling scan_flush", scan.external_id)
celery_brain.scan_flush(scan.external_id)
def set_launched(scanid, scan_report_dict):
""" set status launched for scan
:param scanid: id returned by scan_new
:param scanreport: scan details output by brain
:return: None
:raise: IrmaDatabaseError
"""
with session_transaction() as session:
log.info("scanid: %s is now launched", format(scanid))
scan = Scan.load_from_ext_id(scanid, session=session)
if scan.status == IrmaScanStatus.uploaded:
scan.set_status(IrmaScanStatus.launched)
session.commit()
def get(scanid, db):
""" Retrieve information for a specific scan
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def get(scanid, db):
""" Retrieve information for a specific scan
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def cancel(scanid, db):
""" Cancel a scan.
The request should be performed using a POST request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
scan_ctrl.cancel(scan, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def cancel(scanid, db):
""" Cancel a scan.
The request should be performed using a POST request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
scan_ctrl.cancel(scan, db)
response.content_type = "application/json; charset=UTF-8"
return scan_schema.dumps(scan).data
except Exception as e:
log.exception(e)
process_error(e)
def get_results(scanid, db):
""" Retrieve results for a scan. Results are the same as in the get()
method, i.e. a summary for each scanned files.
The request should be performed using a GET request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
file_web_schema = FileWebSchema_v1(exclude=('probe_results',
'file_infos'))
response.content_type = "application/json; charset=UTF-8"
return file_web_schema.dumps(scan.files_web, many=True).data
except Exception as e:
log.exception(e)
process_error(e)
def get_results(scanid, db):
""" Retrieve results for a scan. Results are the same as in the get()
method, i.e. a summary for each scanned files.
The request should be performed using a GET request method.
"""
try:
log.debug("scanid: %s", scanid)
validate_id(scanid)
scan = Scan.load_from_ext_id(scanid, db)
file_web_schema = FileWebSchema_v1(exclude=('probe_results',
'file_infos'))
response.content_type = "application/json; charset=UTF-8"
return file_web_schema.dumps(scan.files_web, many=True).data
except Exception as e:
log.exception(e)
process_error(e)
def handle_output_files(scanid, parent_file_hash, probe, result):
with session_transaction() as session:
scan = Scan.load_from_ext_id(scanid, session=session)
uploaded_files = result.get('uploaded_files', None)
if uploaded_files is None or not scan.resubmit_files:
log.debug("scanid: %s Nothing to resubmit or resubmit disabled",
scanid)
return
log.info("scanid: %s appending new uploaded files %s", scanid,
uploaded_files.keys())
parent_file = File.load_from_sha256(parent_file_hash, session)
# filter already present file in current scan
hash_uploaded = [f.sha256 for f in scan.files]
new_fws = _append_new_files_to_scan(scan, uploaded_files, session)
for fw in new_fws:
parent_file.children.append(fw)
_resubmit_files(scan, parent_file, new_fws, hash_uploaded, session)
def setUp(self):
self.date = "date"
self.ip = "ip"
self.scan = Scan(self.date, self.ip)
class TestScan(TestCase):
def setUp(self):
self.date = "date"
self.ip = "ip"
self.scan = Scan(self.date, self.ip)
def tearDown(self):
del self.scan
def test001_finished_not_uploaded(self):
status = choice([IrmaScanStatus.empty,
IrmaScanStatus.ready])
self.scan.set_status(status)
self.assertFalse(self.scan.finished())
def test002_finished_launched_not_finished(self):
a, b = MagicMock(), MagicMock()
a.doc = "something"
b.doc = None
fw = MagicMock()
fw.probe_results = [a, b]
self.scan.files_web = [fw]
self.scan.set_status(IrmaScanStatus.launched)
self.assertFalse(self.scan.finished())
def test003_finished_launched_finished(self):
a, b = MagicMock(), MagicMock()
a.doc = "something"
b.doc = "anotherthing"
fw = MagicMock()
fw.probe_results = [a, b]
self.scan.files_web = [fw]
self.scan.set_status(IrmaScanStatus.launched)
self.assertTrue(self.scan.finished())
def test004_finished_finished(self):
self.scan.set_status(IrmaScanStatus.launched)
res = self.scan.finished()
self.assertTrue(res)
def test005_probes_total(self):
fw1, fw2 = MagicMock(), MagicMock()
pt1, pt2 = randint(0, 20), randint(0, 20)
fw1.probes_total = pt1
fw2.probes_total = pt2
self.scan.files_web = [fw1, fw2]
self.assertEquals(self.scan.probes_total, pt1 + pt2)
def test006_probes_finished(self):
fw1, fw2 = MagicMock(), MagicMock()
pf1, pf2 = randint(0, 20), randint(0, 20)
fw1.probes_finished = pf1
fw2.probes_finished = pf2
self.scan.files_web = [fw1, fw2]
self.assertEquals(self.scan.probes_finished, pf1 + pf2)
def test007_files(self):
fw = MagicMock()
self.scan.files_web = [fw]
self.assertEqual(self.scan.files, [fw.file])
def test008_set_status(self):
with self.assertRaises(IrmaCoreError):
self.scan.set_status("whatever")
def test009_fileweb_by_sha256(self):
fw = MagicMock()
sha256 = "whatever"
fw.file.sha256 = sha256
self.scan.files_web = [fw]
self.assertEqual(self.scan.get_filewebs_by_sha256(sha256), [fw])
def test010_query_find_by_filesha256(self):
m_session = MagicMock()
sha256 = "whatever"
Scan.query_find_by_filesha256(sha256, m_session)
m_filter = m_session.query(Scan).join().join().filter
m_filter.is_called_once_with(File.sha256 == sha256)
def test011_load_from_ext_id(self):
m_session = MagicMock()
m_query = MagicMock()
m_session.query.return_value = m_query
external_id = "whatever"
Scan.load_from_ext_id(external_id, m_session)
m_filter = m_query.filter
m_filter.assert_called_once()
def test012_load_from_ext_id_none(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().filter
m_filter.side_effect = NoResultFound()
with self.assertRaises(IrmaDatabaseResultNotFound):
Scan.load_from_ext_id(external_id, m_session)
def test013_load_from_ext_id_multiple(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().filter
m_filter.side_effect = MultipleResultsFound()
with self.assertRaises(IrmaDatabaseError):
Scan.load_from_ext_id(external_id, m_session)
def test014_set_get_probelist(self):
probelist = ["probe1", "probe2"]
self.scan.set_probelist(probelist)
self.assertItemsEqual(self.scan.get_probelist(), probelist)
def test015_finished(self):
self.scan.set_status(IrmaScanStatus.finished)
self.assertTrue(self.scan.finished())
def test010_query_find_by_filesha256(self):
m_session = MagicMock()
sha256 = "whatever"
Scan.query_find_by_filesha256(sha256, m_session)
m_filter = m_session.query(Scan).join().join().filter
m_filter.is_called_once_with(File.sha256 == sha256)
def setUp(self):
self.date = "date"
self.ip = "ip"
self.scan = Scan(self.date, self.ip)
class TestScan(TestCase):
def setUp(self):
self.date = "date"
self.ip = "ip"
self.scan = Scan(self.date, self.ip)
def tearDown(self):
del self.scan
def test001_finished_not_uploaded(self):
status = choice([IrmaScanStatus.empty, IrmaScanStatus.ready])
self.scan.set_status(status)
self.assertFalse(self.scan.finished())
def test002_finished_launched_not_finished(self):
a, b = MagicMock(), MagicMock()
a.doc = "something"
b.doc = None
fw = MagicMock()
fw.probe_results = [a, b]
self.scan.files_web = [fw]
self.scan.set_status(IrmaScanStatus.launched)
self.assertFalse(self.scan.finished())
def test003_finished_launched_finished(self):
a, b = MagicMock(), MagicMock()
a.doc = "something"
b.doc = "anotherthing"
fw = MagicMock()
fw.probe_results = [a, b]
self.scan.files_web = [fw]
self.scan.set_status(IrmaScanStatus.launched)
self.assertTrue(self.scan.finished())
def test004_finished_finished(self):
self.scan.set_status(IrmaScanStatus.launched)
res = self.scan.finished()
self.assertTrue(res)
def test005_probes_total(self):
fw1, fw2 = MagicMock(), MagicMock()
pt1, pt2 = randint(0, 20), randint(0, 20)
fw1.probes_total = pt1
fw2.probes_total = pt2
self.scan.files_web = [fw1, fw2]
self.assertEquals(self.scan.probes_total, pt1 + pt2)
def test006_probes_finished(self):
fw1, fw2 = MagicMock(), MagicMock()
pf1, pf2 = randint(0, 20), randint(0, 20)
fw1.probes_finished = pf1
fw2.probes_finished = pf2
self.scan.files_web = [fw1, fw2]
self.assertEquals(self.scan.probes_finished, pf1 + pf2)
def test007_files(self):
fw = MagicMock()
self.scan.files_web = [fw]
self.assertEqual(self.scan.files, [fw.file])
def test008_set_status(self):
with self.assertRaises(IrmaCoreError):
self.scan.set_status("whatever")
def test009_fileweb_by_sha256(self):
fw = MagicMock()
sha256 = "whatever"
fw.file.sha256 = sha256
self.scan.files_web = [fw]
self.assertEqual(self.scan.get_filewebs_by_sha256(sha256), [fw])
def test010_query_find_by_filesha256(self):
m_session = MagicMock()
sha256 = "whatever"
Scan.query_find_by_filesha256(sha256, m_session)
m_filter = m_session.query(Scan).join().join().filter
m_filter.is_called_once_with(File.sha256 == sha256)
def test011_load_from_ext_id(self):
m_session = MagicMock()
m_query = MagicMock()
m_session.query.return_value = m_query
external_id = "whatever"
Scan.load_from_ext_id(external_id, m_session)
m_filter = m_query.options().filter
m_filter.assert_called_once()
def test012_load_from_ext_id_none(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().options().filter
m_filter.side_effect = NoResultFound()
with self.assertRaises(IrmaDatabaseResultNotFound):
Scan.load_from_ext_id(external_id, m_session)
def test013_load_from_ext_id_multiple(self):
m_session = MagicMock()
external_id = "whatever"
m_filter = m_session.query().options().filter
m_filter.side_effect = MultipleResultsFound()
with self.assertRaises(IrmaDatabaseError):
Scan.load_from_ext_id(external_id, m_session)
def test014_set_get_probelist(self):
probelist = ["probe1", "probe2"]
self.scan.set_probelist(probelist)
self.assertItemsEqual(self.scan.get_probelist(), probelist)
def test015_finished(self):
self.scan.set_status(IrmaScanStatus.finished)
self.assertTrue(self.scan.finished())
def test010_query_find_by_filesha256(self):
m_session = MagicMock()
sha256 = "whatever"
Scan.query_find_by_filesha256(sha256, m_session)
m_filter = m_session.query(Scan).join().join().filter
m_filter.is_called_once_with(File.sha256 == sha256)
