Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes #114 #124

Merged
merged 1 commit into from
Apr 27, 2018
Merged

Fixes #114 #124

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 25 additions & 12 deletions utils/celery_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ def multiscanner_celery(file_, original_filename, task_id, file_hash, metadata,
multiscanner_celery.delay(full_path, original_filename, task_id,
hashed_filename, metadata, config, module_list)
'''

# Initialize the connection to the task DB
db.init_db()

Expand Down Expand Up @@ -163,15 +164,20 @@ def multiscanner_celery(file_, original_filename, task_id, file_hash, metadata,
# Count number of modules enabled out of total possible
# and add it to the Scan Metadata
total_enabled = 0
total_modules = 0
for key in full_conf:
if key == 'main':
continue
sub_conf[key] = {}
sub_conf[key]['ENABLED'] = full_conf[key]['ENABLED']
total_modules += 1
if sub_conf[key]['ENABLED'] is True:
total_enabled += 1
total_modules = len(full_conf.keys())

# Get the count of modules enabled from the module_list
# if it exists, else count via the config
if module_list:
total_enabled = len(module_list)
else:
for key in full_conf:
if key == 'main':
continue
sub_conf[key] = {}
sub_conf[key]['ENABLED'] = full_conf[key]['ENABLED']
if sub_conf[key]['ENABLED'] is True:
total_enabled += 1

results[file_]['Scan Metadata'] = metadata
results[file_]['Scan Metadata']['Worker Node'] = gethostname()
Expand All @@ -188,16 +194,23 @@ def multiscanner_celery(file_, original_filename, task_id, file_hash, metadata,
results[original_filename] = results[file_]
del results[file_]

# Save the reports to storage
storage_ids = storage_handler.store(results, wait=False)
storage_handler.close()

# Only need to raise ValueError here,
# Further cleanup will be handled by the on_failure method
# of MultiScannerTask
if not storage_ids:
raise ValueError('Report failed to index')

# Update the task DB to reflect that the task is done
db.update_task(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not exactly related to this PR, but while we're on the subject...should we be doing any checks to see if db.update_task was successful?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good point. I added an issue for this (#129), since I think it will need a little more of a reworking of the sql_driver.py. I think this sort of error handling should be handled by the ORM layer, with an exception or w/e raised to the caller. Thoughts?

task_id=task_id,
task_status='Complete',
timestamp=scan_time,
)

# Save the reports to storage
storage_handler.store(results, wait=False)
storage_handler.close()
logger.info('Completed Task #{}'.format(task_id))

return results
Expand Down