Skip to content

Commit

Permalink
Merge pull request #160 from emmanvg/files-endpoint-fix
Browse files Browse the repository at this point in the history 
Fixes for API Endpoint (#159)
  • Loading branch information
@ptcNOP
ptcNOP committed Nov 19, 2018
2 parents dedc48c + 20f8788 commit c603fae
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 21 deletions.
46 changes: 44 additions & 2 deletions tests/test_api.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -283,7 +283,7 @@ def test_edit_notes(self, mock_handler):

args, kwargs = mock_handler.edit_note.call_args_list[0]
self.assertEqual(args[0], '114d70ba7d04c76d8c217c970f99682025c89b1a6ffe91eb9045653b4b954eb9')
self.assertEqual(args[1], '1')
self.assertEqual(args[1], 1)
self.assertEqual(args[2], 'bar')

@mock.patch('api.handler')
Expand All @@ -292,4 +292,46 @@ def test_remove_notes(self, mock_handler):

args, kwargs = mock_handler.delete_note.call_args_list[0]
self.assertEqual(args[0], '114d70ba7d04c76d8c217c970f99682025c89b1a6ffe91eb9045653b4b954eb9')
self.assertEqual(args[1], '1')
self.assertEqual(args[1], 1)


class TestSHA256DownloadSampleCase(APITestCase):
def setUp(self):
super(self.__class__, self).setUp()
# populate the DB w/ a task
post_file(self.app)
self.sql_db.update_task(
task_id=1,
task_status='Complete',
)

@mock.patch('api.db')
@mock.patch('api.handler')
def test_malformed_request(self, mock_handler, mock_db):
resp = self.app.get('/api/v1/files/..\opt\multiscanner\web_config.ini')

self.assertEqual(resp.status_code, api.HTTP_BAD_REQUEST)

@mock.patch('api.db')
@mock.patch('api.handler')
def test_other_hash(self, mock_handler, mock_db):
# using MD5 instead of SHA256
resp = self.app.get('/api/v1/files/96b47da202ddba8d7a6b91fecbf89a41')

self.assertEqual(resp.status_code, api.HTTP_BAD_REQUEST)

@mock.patch('api.db')
@mock.patch('api.handler')
def test_file_download_raw(self, mock_handler, mock_db):
expected_response = b'my file contents'
resp = self.app.get('/api/v1/files/114d70ba7d04c76d8c217c970f99682025c89b1a6ffe91eb9045653b4b954eb9?raw=t')

self.assertEqual(resp.status_code, api.HTTP_OK)
self.assertEqual(resp.get_data(), expected_response)

@mock.patch('api.db')
@mock.patch('api.handler')
def test_file_not_found(self, mock_handler, mock_db):
resp = self.app.get('/api/v1/files/26d11f0ea5cc77a59b6e47deee859440f26d2d14440beb712dbac8550d35ef1f?raw=t')

self.assertEqual(resp.status_code, api.HTTP_NOT_FOUND)
42 changes: 23 additions & 19 deletions utils/api.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import multiprocessing
import os
import queue
import re
import shutil
import subprocess
import sys
Expand All @@ -56,7 +57,7 @@

import rarfile
import requests
from flask import Flask, abort, jsonify, make_response, request
from flask import Flask, abort, jsonify, make_response, request, safe_join
from flask.json import JSONEncoder
from flask_cors import CORS
from jinja2 import Markup
Expand Down Expand Up @@ -541,7 +542,7 @@ def create_task():
)


@app.route('/api/v1/tasks/<task_id>/report', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/report', methods=['GET'])
def get_report(task_id):
'''
Return a JSON dictionary corresponding
Expand Down Expand Up @@ -572,7 +573,7 @@ def _pre_process(report_dict={}):
executed on report_dict.
'''

# pop unecessary keys
# pop unnecessary keys
if report_dict.get('Report', {}).get('ssdeep', {}):
for k in ['chunksize', 'chunk', 'double_chunk']:
try:
Expand Down Expand Up @@ -628,7 +629,7 @@ def _linkify(s, url, new_tab=True):
s=s)


@app.route('/api/v1/tasks/<task_id>/file', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/file', methods=['GET'])
def files_get_task(task_id):
# try to get report dict
report_dict, success = get_report_dict(task_id)
Expand All @@ -645,7 +646,7 @@ def files_get_task(task_id):
return jsonify({'Error': 'sha256 not in report!'})


@app.route('/api/v1/tasks/<task_id>/maec', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/maec', methods=['GET'])
def get_maec_report(task_id):
# try to get report dict
report_dict, success = get_report_dict(task_id)
Expand Down Expand Up @@ -695,7 +696,7 @@ def taglist():
return jsonify({'Tags': response})


@app.route('/api/v1/tasks/<task_id>/tags', methods=['POST', 'DELETE'])
@app.route('/api/v1/tasks/<int:task_id>/tags', methods=['POST', 'DELETE'])
def tags(task_id):
'''
Add/Remove the specified tag to the specified task.
Expand All @@ -719,7 +720,7 @@ def tags(task_id):
return jsonify({'Message': 'Tag Removed'})


@app.route('/api/v1/tasks/<task_id>/notes', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/notes', methods=['GET'])
def get_notes(task_id):
'''
Get one or more analyst notes/comments associated with the specified task.
Expand Down Expand Up @@ -749,7 +750,7 @@ def get_notes(task_id):
return jsonify(response)


@app.route('/api/v1/tasks/<task_id>/notes', methods=['POST'])
@app.route('/api/v1/tasks/<int:task_id>/notes', methods=['POST'])
def add_note(task_id):
'''
Add an analyst note/comment to the specified task.
Expand All @@ -764,7 +765,7 @@ def add_note(task_id):
return jsonify(response)


@app.route('/api/v1/tasks/<task_id>/notes/<note_id>', methods=['PUT', 'DELETE'])
@app.route('/api/v1/tasks/<int:task_id>/notes/<int:note_id>', methods=['PUT', 'DELETE'])
def edit_note(task_id, note_id):
'''
Modify/remove the specified analyst note/comment.
Expand All @@ -784,7 +785,7 @@ def edit_note(task_id, note_id):
return jsonify(response)


@app.route('/api/v1/files/<sha256>', methods=['GET'])
@app.route('/api/v1/files/<string:sha256>', methods=['GET'])
# get raw file - /api/v1/files/get/<sha256>?raw=true
def files_get_sha256(sha256):
'''
Expand All @@ -793,18 +794,21 @@ def files_get_sha256(sha256):
# is there a robust way to just get this as a bool?
raw = request.args.get('raw', default='False', type=str)

return files_get_sha256_helper(sha256, raw)
if re.match(r'^[a-fA-F0-9]{64}$', sha256):
return files_get_sha256_helper(sha256, raw)
else:
return abort(HTTP_BAD_REQUEST)


def files_get_sha256_helper(sha256, raw=None):
'''
Returns binary from storage. Defaults to password protected zipfile.
'''
file_path = os.path.join(api_config['api']['upload_folder'], sha256)
file_path = safe_join(api_config['api']['upload_folder'], sha256)
if not os.path.exists(file_path):
abort(HTTP_NOT_FOUND)

with open(file_path, "rb") as fh:
with open(file_path, 'rb') as fh:
fh_content = fh.read()

raw = raw[0].lower()
Expand All @@ -816,13 +820,13 @@ def files_get_sha256_helper(sha256, raw=None):
else:
# ref: https://github.com/crits/crits/crits/core/data_tools.py#L122
rawname = sha256 + '.bin'
with open(os.path.join('/tmp/', rawname), 'wb') as raw_fh:
with open(safe_join('/tmp/', rawname), 'wb') as raw_fh:
raw_fh.write(fh_content)

zipname = sha256 + '.zip'
args = ['/usr/bin/zip', '-j',
os.path.join('/tmp', zipname),
os.path.join('/tmp', rawname),
safe_join('/tmp', zipname),
safe_join('/tmp', rawname),
'-P', 'infected']
proc = subprocess.Popen(args)
wait_seconds = 30
Expand All @@ -836,7 +840,7 @@ def files_get_sha256_helper(sha256, raw=None):
proc.terminate()
return make_response(jsonify({'Error': 'Process timed out'}))
else:
with open(os.path.join('/tmp', zipname), 'rb') as zip_fh:
with open(safe_join('/tmp', zipname), 'rb') as zip_fh:
zip_data = zip_fh.read()
if len(zip_data) == 0:
return make_response(jsonify({'Error': 'Zip file empty'}))
Expand Down Expand Up @@ -881,7 +885,7 @@ def run_ssdeep_group():
HTTP_BAD_REQUEST)


@app.route('/api/v1/tasks/<task_id>/pdf', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/pdf', methods=['GET'])
def get_pdf_report(task_id):
'''
Generates a PDF version of a JSON report.
Expand All @@ -898,7 +902,7 @@ def get_pdf_report(task_id):
return response


@app.route('/api/v1/tasks/<task_id>/stix2', methods=['GET'])
@app.route('/api/v1/tasks/<int:task_id>/stix2', methods=['GET'])
def get_stix2_bundle_from_report(task_id):
'''
Generates a STIX2 Bundle with indicators generated of a JSON report.
Expand Down

0 comments on commit c603fae

Please sign in to comment.