def scan(frontend_scanid, scan_request_dict):
try:
with session_transaction() as session:
log.debug("scanid: %s received %s", frontend_scanid,
scan_request_dict)
user = User.get_by_rmqvhost(session)
scan_request = IrmaScanRequest(scan_request_dict)
scan = scan_ctrl.new(frontend_scanid, user, scan_request.nb_files,
session)
available_probelist = probe_ctrl.get_list(session)
# Now, create one subtask per file to
# scan per probe
new_jobs = []
for filehash in scan_request.filehashes():
probelist = scan_request.get_probelist(filehash)
scan_ctrl.check_probelist(scan, probelist, available_probelist,
session)
for probe in probelist:
j = Job(scan.id, filehash, probe)
session.add(j)
new_jobs.append(j)
session.commit()
scan_ctrl.launch(scan, new_jobs, session)
celery_frontend.scan_launched(scan.scan_id, scan_request.to_dict())
log.info(
"scanid %s: %d file(s) received / %d active probe(s) / "
"%d job(s) launched", scan.scan_id, scan.nb_files,
len(available_probelist), len(scan.jobs))
except Exception as e:
log.exception(e)
raise
def scan(file, probelist, frontend_scan):
try:
with session_transaction() as session:
log.debug("scan_id: %s fileweb_id: %s received %s", frontend_scan,
file, probelist)
user = User.get_by_rmqvhost(session)
scan = scan_ctrl.new(frontend_scan, user, session)
# Now, create one subtask per file to scan per probe
new_jobs = []
for probe in probelist:
if probe in probe_ctrl.available_probes:
j = Job(scan.id, file, probe)
session.add(j)
new_jobs.append(j)
else:
# send an error message to not stuck the scan
# One probe asked is no more present
res = probe_ctrl.create_error_results(probe,
"missing probe",
session)
celery_frontend.scan_result(file,
probe,
res)
session.commit()
scan_ctrl.launch(scan, new_jobs, session)
log.info("scan_id %s: file %s received / %d active probe(s) / "
"%d job(s) launched",
scan.scan_id,
file,
len(probe_ctrl.available_probes),
len(scan.jobs))
except Exception as e:
log.exception(e)
raise
def scan(file, probelist, frontend_scan):
try:
with session_transaction() as session:
log.debug("scan_id: %s fileweb_id: %s received %s", frontend_scan,
file, probelist)
user = User.get_by_rmqvhost(session)
scan = scan_ctrl.new(frontend_scan, user, session)
# Now, create one subtask per file to scan per probe
new_jobs = []
for probe in probelist:
if probe in probe_ctrl.available_probes:
j = Job(scan.id, file, probe)
session.add(j)
new_jobs.append(j)
else:
# send an error message to not stuck the scan
# One probe asked is no more present
res = probe_ctrl.create_error_results(
probe, "missing probe", session)
celery_frontend.scan_result(file, probe, res)
session.commit()
scan_ctrl.launch(scan, new_jobs, session)
log.info(
"scan_id %s: file %s received / %d active probe(s) / "
"%d job(s) launched", scan.scan_id, file,
len(probe_ctrl.available_probes), len(scan.jobs))
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
raise
def scan(frontend_scanid, scan_request_dict):
try:
with session_transaction() as session:
log.debug("scanid: %s received %s", frontend_scanid,
scan_request_dict)
user = User.get_by_rmqvhost(session)
scan_request = IrmaScanRequest(scan_request_dict)
scan = scan_ctrl.new(frontend_scanid, user, scan_request.nb_files,
session)
available_probelist = probe_ctrl.get_list(session)
# Now, create one subtask per file to
# scan per probe
new_jobs = []
for filehash in scan_request.filehashes():
probelist = scan_request.get_probelist(filehash)
scan_ctrl.check_probelist(scan, probelist,
available_probelist, session)
for probe in probelist:
j = Job(scan.id, filehash, probe)
session.add(j)
new_jobs.append(j)
session.commit()
scan_ctrl.launch(scan, new_jobs, session)
celery_frontend.scan_launched(scan.scan_id,
scan_request.to_dict())
log.info("scanid %s: %d file(s) received / %d active probe(s) / "
"%d job(s) launched",
scan.scan_id,
scan.nb_files,
len(available_probelist),
len(scan.jobs))
except Exception as e:
log.exception(e)
raise
def get_quota(user_id):
with session_query() as session:
user = User.load(user_id, session)
quota = user.quota
remaining = None
if quota is not None:
remaining = user.remaining_quota(session)
return (remaining, quota)
def setUp(self):
self.user_name = "test_user"
self.user_rmqvhost = "test_vhost"
self.user_ftpuser = "******"
self.user_quota = 100
self.scanid = UUID.generate()
try:
self.userid = user_ctrl.get_userid(self.user_rmqvhost)
except IrmaDatabaseResultNotFound:
# user does not exist create one
with session_transaction() as session:
user = User(self.user_name,
self.user_rmqvhost,
self.user_ftpuser,
self.user_quota)
user.save(session)
session.commit()
self.userid = user.id
def scan_flush(frontend_scan_id):
try:
log.info("scan_id %s: scan flush requested", frontend_scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(e)
return
def scan_flush(scan_id):
try:
log.info("scan %s: flush requested", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(e)
return
def scan_cancel(frontend_scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(e)
return IrmaTaskReturn.error("cancel error on brain")
def scan_flush(scan_id):
try:
log.info("scan %s: flush requested", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return
def scan_cancel(frontend_scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(e)
return IrmaTaskReturn.error("cancel error on brain")
def files_flush(files, scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
ftpuser = scan.user.ftpuser
log.debug("Flushing files %s", files)
ftp_ctrl.flush(ftpuser, files)
except Exception as e:
log.exception(e)
return
def setUp(self):
self.session = MagicMock()
self.name = "user"
self.rmqvhost = "vhost"
self.ftpuser = "******"
self.user = User(self.name, self.rmqvhost, self.ftpuser)
self.frontend_scanid = "frontend_scanid"
self.nb_files = randint(50, 100)
self.scan = Scan(frontend_scanid=self.frontend_scanid,
user_id=self.user.id)
self.scan.user = self.user
def files_flush(files, scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
ftpuser = scan.user.ftpuser
log.debug("Flushing files %s", files)
ftp_ctrl.flush(ftpuser, files)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return
def scan_cancel(scan_id):
try:
log.info("scan %s: cancel", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return IrmaTaskReturn.error("cancel error on brain")
def test_user_quota_update(self):
# test update quota
(before, _) = user_ctrl.get_quota(self.userid)
scanid = scan_ctrl.new(self.scanid, self.userid, 2)
for i in xrange(0, 10):
for probe in ['probe1', 'probe2']:
job_ctrl.new(scanid, "file-{0}".format(i), probe, 1)
scan_ctrl.launched(scanid)
(after, _) = user_ctrl.get_quota(self.userid)
self.assertEqual(before - after, 20)
with session_query() as session:
user = User.load(self.userid, session)
scan_ids = [scan.id for scan in user.scans]
self.assertNotEqual(scan_ids, [])
self.assertTrue(scanid in scan_ids)
if len(sys.argv) not in (4, 5):
print("usage: {0} <username> <rmqvhost> <ftpuser> [quota]\n"
" with <username> a string\n"
" <rmqvhost> the rmqvhost used for the frontend\n"
" <ftpuser> the ftpuser used by the frontend\n"
" [quota] the number of file scan quota\n"
"example: {0} test1 mqfrontend frontend"
"".format(sys.argv[0]))
sys.exit(1)
name = sys.argv[1]
rmqvhost = sys.argv[2]
ftpuser = sys.argv[3]
# quota is in number of files (0 means disabled)
quota = int(sys.argv[4]) if len(sys.argv) == 5 else 0
with session_transaction() as session:
try:
user = User.get_by_rmqvhost(rmqvhost, session)
print("rmqvhost {0} is already assigned to user {1}. "
"Updating with new parameters.".format(user.name, user.rmqvhost))
user = user.load(user.id, session)
user.name = name
user.ftpuser = ftpuser
user.quota = quota
user.update(['name', 'ftpuser', 'quota'], session)
except IrmaDatabaseResultNotFound:
user = User(name=name, rmqvhost=rmqvhost, ftpuser=ftpuser, quota=quota)
user.save(session)
print("SQL directory does not exist {0}" "..creating".format(dirname))
os.makedirs(dirname)
os.chmod(dirname, 0o777)
elif not (os.path.isdir(dirname)):
print("Error. SQL directory is a not a dir {0}" "".format(dirname))
raise IrmaDatabaseError("Can not create Brain database dir")
if not os.path.exists(db_name):
# touch like method to create a rw-rw-rw- file for db
open(db_name, 'a').close()
os.chmod(db_name, 0o666)
# Retrieve database informations
url = config.get_sql_url()
engine = create_engine(url, echo=config.sql_debug_enabled())
# and create Database in case
Base.metadata.create_all(engine)
with session_transaction() as session:
try:
user = User.get_by_rmqvhost(session, rmqvhost=rmqvhost)
print("rmqvhost {0} is already assigned to user {1}. "
"Updating with new parameters.".format(user.name, user.rmqvhost))
user = user.load(user.id, session)
user.name = name
user.ftpuser = ftpuser
user.update(['name', 'ftpuser'], session)
except IrmaDatabaseResultNotFound:
user = User(name=name, rmqvhost=rmqvhost, ftpuser=ftpuser)
user.save(session)
os.makedirs(dirname)
os.chmod(dirname, 0777)
elif not (os.path.isdir(dirname)):
print("Error. SQL directory is a not a dir {0}"
"".format(dirname))
raise IrmaDatabaseError("Can not create Brain database dir")
if not os.path.exists(db_name):
# touch like method to create a rw-rw-rw- file for db
open(db_name, 'a').close()
os.chmod(db_name, 0666)
# Retrieve database informations
url = config.get_sql_url()
engine = create_engine(url, echo=config.sql_debug_enabled())
# and create Database in case
Base.metadata.create_all(engine)
with session_transaction() as session:
try:
user = User.get_by_rmqvhost(session, rmqvhost=rmqvhost)
print("rmqvhost {0} is already assigned to user {1}. "
"Updating with new parameters.".format(user.name, user.rmqvhost))
user = user.load(user.id, session)
user.name = name
user.ftpuser = ftpuser
user.update(['name', 'ftpuser'], session)
except IrmaDatabaseResultNotFound:
user = User(name=name, rmqvhost=rmqvhost, ftpuser=ftpuser)
user.save(session)
def test003_get_by_rmqvhost_not_found(self):
self.session.query.side_effect = NoResultFound
with self.assertRaises(IrmaDatabaseResultNotFound):
User.get_by_rmqvhost(self.session, rmqvhost="whatever")
print("SQL directory does not exist {0}"
"..creating".format(dirname))
os.makedirs(dirname)
os.chmod(dirname, 0o777)
elif not (os.path.isdir(dirname)):
print("Error. SQL directory is a not a dir {0}"
"".format(dirname))
raise IrmaDatabaseError("Can not create Brain database dir")
if not os.path.exists(db_name):
# touch like method to create a rw-rw-rw- file for db
open(db_name, 'a').close()
os.chmod(db_name, 0o666)
# Retrieve database informations
engine = create_engine(config.sqldb.url, echo=config.sql_debug_enabled())
# and create Database in case
Base.metadata.create_all(engine)
with session_transaction() as session:
try:
user = User.get_by_rmqvhost(session, rmqvhost=rmqvhost)
print("rmqvhost {0} is already assigned to user {1}. "
"Updating with new parameters.".format(user.name, user.rmqvhost))
session.query(User)\
.filter_by(id=user.id)\
.update({'name': name, 'ftpuser': ftpuser})
except IrmaDatabaseResultNotFound:
user = User(name=name, rmqvhost=rmqvhost, ftpuser=ftpuser)
session.add(user)
def test001___init__(self):
user = User(self.name, self.rmqvhost, self.ftpuser)
self.assertEqual(user.name, self.name)
self.assertEqual(user.rmqvhost, self.rmqvhost)
self.assertEqual(user.ftpuser, self.ftpuser)
def test002_get_by_rmqvhost(self):
User.get_by_rmqvhost(self.session)
self.session.query.assert_called_once_with(User)
m_filter = self.session.query(User).filter
m_filter.assert_called_once()
m_filter().one.assert_called_once()
def get_ftpuser(user_id):
with session_query() as session:
user = User.load(user_id, session)
return user.ftpuser
def test004_get_by_rmqvhost_multiple_found(self):
self.session.query.side_effect = MultipleResultsFound
with self.assertRaises(IrmaDatabaseError):
User.get_by_rmqvhost(self.session, rmqvhost="whatever")
def get_userid(rmqvhost):
with session_query() as session:
user = User.get_by_rmqvhost(rmqvhost, session)
return user.id
