def test_heuristics_valid():
heuristic_list = [random_model_obj(Heuristic) for _ in range(4)]
heuristics = {x.heur_id: x for x in heuristic_list}
attack_ids = list(
set([
random.choice(list(attack_map.keys()))
for _ in range(random.randint(1, 3))
]))
signatures = {}
score_map = {}
for x in range(random.randint(2, 4)):
name = get_random_word()
if x >= 2:
score_map[name] = random.randint(10, 100)
signatures[name] = random.randint(1, 3)
service_heur = dict(heur_id=random.choice(list(heuristics.keys())),
score=0,
attack_ids=attack_ids,
signatures=signatures,
frequency=0,
score_map=score_map)
result_heur = service_heuristic_to_result_heuristic(
deepcopy(service_heur), heuristics)
assert result_heur is not None
assert service_heur['heur_id'] == result_heur['heur_id']
assert service_heur['score'] != result_heur['score']
for attack in result_heur['attack']:
assert attack['attack_id'] in attack_ids
for signature in result_heur['signature']:
assert signature['name'] in signatures
assert signature['frequency'] == signatures[signature['name']]
def test_heuristics_valid():
heuristic_list = [random_model_obj(Heuristic) for _ in range(4)]
heuristics = {x.heur_id: x for x in heuristic_list}
software_ids = list(set([random.choice(list(software_map.keys())) for _ in range(random.randint(1, 3))]))
attack_ids = list(set([random.choice(list(attack_map.keys())) for _ in range(random.randint(1, 3))]))
attack_ids_to_fetch_details_for = attack_ids[:]
for software_id in software_ids:
software_attack_ids = software_map[software_id]["attack_ids"]
for software_attack_id in software_attack_ids:
if software_attack_id in attack_map and software_attack_id not in attack_ids_to_fetch_details_for:
attack_ids_to_fetch_details_for.append(software_attack_id)
else:
print(f"Invalid related attack_id '{software_attack_id}' for software '{software_id}'. Ignoring it.")
attack_id_details = {attack_id: {"pattern": attack_map[attack_id]["name"], "categories": attack_map[attack_id]["categories"]} for attack_id in attack_ids_to_fetch_details_for}
attack_ids.extend(software_ids)
signatures = {}
score_map = {}
for x in range(random.randint(2, 4)):
name = get_random_word()
if x >= 2:
score_map[name] = random.randint(10, 100)
signatures[name] = random.randint(1, 3)
service_heur = dict(
heur_id=random.choice(list(heuristics.keys())),
score=0,
attack_ids=attack_ids,
signatures=signatures,
frequency=0,
score_map=score_map
)
result_heur = service_heuristic_to_result_heuristic(deepcopy(service_heur), heuristics)
assert result_heur is not None
assert service_heur['heur_id'] == result_heur['heur_id']
assert service_heur['score'] != result_heur['score']
for attack in result_heur['attack']:
attack_id = attack['attack_id']
assert attack_id in attack_ids_to_fetch_details_for
assert attack['pattern'] == attack_id_details[attack_id]['pattern']
assert attack['categories'] == attack_id_details[attack_id]['categories']
for signature in result_heur['signature']:
assert signature['name'] in signatures
assert signature['frequency'] == signatures[signature['name']]
def test_heuristics_invalid():
with pytest.raises(InvalidHeuristicException):
service_heuristic_to_result_heuristic({'heur_id': "my_id"}, {})
def try_run(self):
try:
self.service_class = load_module_by_path(SERVICE_PATH)
except ValueError:
raise
except Exception:
LOG.error("Could not find service in path. Check your environment variables.")
raise
self.load_service_manifest()
if not os.path.isfile(FILE_PATH):
LOG.info(f"File not found: {FILE_PATH}")
return
self.file_dir = os.path.dirname(FILE_PATH)
# Get filename and working dir
file_name = os.path.basename(FILE_PATH)
working_dir = os.path.join(self.file_dir, f'{os.path.basename(FILE_PATH)}_{SERVICE_NAME.lower()}')
# Start service
self.service.start_service()
# Identify the file
file_info = identify.fileinfo(FILE_PATH)
if file_info['type'] == "archive/cart":
# This is a CART file, uncart it and recreate the file info object
original_temp = os.path.join(tempfile.gettempdir(), file_info['sha256'])
with open(FILE_PATH, 'rb') as ifile, open(original_temp, 'wb') as ofile:
unpack_stream(ifile, ofile)
file_info = identify.fileinfo(original_temp)
target_file = os.path.join(tempfile.gettempdir(), file_info['sha256'])
shutil.move(original_temp, target_file)
LOG.info(f"File was a CaRT archive, it was un-CaRTed to {target_file} for processing")
else:
# It not a cart, move the file to the right place to be processed
target_file = os.path.join(tempfile.gettempdir(), file_info['sha256'])
shutil.copyfile(FILE_PATH, target_file)
# Create service processing task
service_task = ServiceTask(dict(
sid=get_random_id(),
metadata={},
service_name=SERVICE_NAME,
service_config=self.submission_params,
fileinfo=dict(
magic=file_info['magic'],
md5=file_info['md5'],
mime=file_info['mime'],
sha1=file_info['sha1'],
sha256=file_info['sha256'],
size=file_info['size'],
type=file_info['type'],
),
filename=file_name,
min_classification=forge.get_classification().UNRESTRICTED,
max_files=501, # TODO: get the actual value
ttl=3600,
))
LOG.info(f"Starting task with SID: {service_task.sid}")
# Set the working directory to a directory with same parent as input file
if os.path.isdir(working_dir):
shutil.rmtree(working_dir)
if not os.path.isdir(working_dir):
os.makedirs(os.path.join(working_dir, 'working_directory'))
self.service.handle_task(service_task)
# Move the result.json and extracted/supplementary files to the working directory
source = os.path.join(tempfile.gettempdir(), 'working_directory')
if not os.path.exists(source):
os.makedirs(source)
files = os.listdir(source)
for f in files:
shutil.move(os.path.join(source, f), os.path.join(working_dir, 'working_directory'))
# Cleanup files from the original directory created by the service base
shutil.rmtree(source)
result_json = os.path.join(tempfile.gettempdir(),
f'{service_task.sid}_{service_task.fileinfo.sha256}_result.json')
if not os.path.exists(result_json):
raise Exception("A service error occured and no result json was found.")
# Validate the generated result
with open(result_json, 'r') as fh:
try:
result = json.load(fh)
result.pop('temp_submission_data', None)
for file in result['response']['extracted'] + result['response']['supplementary']:
file.pop('path', None)
# Load heuristics
heuristics = get_heuristics()
# Transform heuristics and calculate score
total_score = 0
for section in result['result']['sections']:
if section['heuristic']:
heur_id = section['heuristic']['heur_id']
try:
section['heuristic'] = service_heuristic_to_result_heuristic(section['heuristic'],
heuristics)
total_score += section['heuristic']['score']
except InvalidHeuristicException:
section['heuristic'] = None
section['heuristic']['name'] = heuristics[heur_id]['name']
result['result']['score'] = total_score
# Add timestamps for creation, archive and expiry
result['created'] = now_as_iso()
result['archive_ts'] = now_as_iso(1 * 24 * 60 * 60)
result['expiry_ts'] = now_as_iso(service_task.ttl * 24 * 60 * 60)
result = Result(result)
# Print the result on console if in debug mode
if args.debug:
f"{SERVICE_NAME.upper()}-RESULT".center(60, '-')
for line in pprint.pformat(result.result.as_primitives()).split('\n'):
LOG.debug(line)
except Exception as e:
LOG.error(f"Invalid result created: {str(e)}")
LOG.info(f"Cleaning up file used for temporary processing: {target_file}")
os.unlink(target_file)
LOG.info(f"Moving {result_json} to the working directory: {working_dir}/result.json")
shutil.move(result_json, os.path.join(working_dir, 'result.json'))
LOG.info(f"Successfully completed task. Output directory: {working_dir}")