def outputs(temp_dir):
with patch.object(Outputs, 'add_collected_command'):
with patch.object(Outputs, 'add_collected_file'):
with patch.object(Outputs, 'add_collected_file_info'):
outputs = Outputs(temp_dir, maxsize=None, sha256=False)
yield outputs
outputs.close()
def test_logging(temp_dir):
# Create an Outputs instance and log a message
output = Outputs(temp_dir, None, False)
logger.info('test log message')
output.close()
# Make sure the log message appears in the output directory
logs = output_file_content(temp_dir, '*-logs.txt')
assert b'test log message' in logs
def test_collect_file_size_filter(temp_dir):
# Create a file that should be collected
test_file = os.path.join(temp_dir, 'test_file.txt')
with open(test_file, 'w') as f:
f.write('content')
# Create a file that should be ignored due to its size
test_big_file = os.path.join(temp_dir, 'test_big_file.txt')
with open(test_big_file, 'w') as f:
f.write('some bigger content')
output = Outputs(temp_dir, '10', False) # Set maximum size to 10 bytes
output.add_collected_file('TestArtifact',
OSFileSystem('/').get_fullpath(test_file))
output.add_collected_file('TestArtifact',
OSFileSystem('/').get_fullpath(test_big_file))
output.close()
zip_content = io.BytesIO(output_file_content(temp_dir, '*-files.zip'))
zipfile = ZipFile(zip_content)
zipped_files = zipfile.namelist()
assert len(zipped_files) == 1
assert zipped_files[0].endswith('test_file.txt')
logs = output_file_content(temp_dir, '*-logs.txt')
assert b"test_big_file.txt' because of its size" in logs
def main(arguments):
try:
locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')
except locale.Error:
pass
output = Outputs(arguments.output, arguments.maxsize, arguments.sha256)
logger.log(PROGRESS, "Loading artifacts ...")
platform = get_operating_system()
collector = Collector(platform)
artifacts_registry = get_artifacts_registry(arguments.library,
arguments.directory)
include_artifacts = resolve_artifact_groups(artifacts_registry,
arguments.include)
exclude_artifacts = resolve_artifact_groups(artifacts_registry,
arguments.exclude)
for artifact_definition, artifact_source in get_artifacts_to_collect(
artifacts_registry, include_artifacts, exclude_artifacts, platform,
arguments.include
or (arguments.directory and not arguments.library)):
collector.register_source(artifact_definition, artifact_source)
collector.collect(output)
def test_collect_pe_file_info(temp_dir, test_pe_file):
output = Outputs(temp_dir, None, False)
output.add_collected_file_info('TestArtifact', test_pe_file)
output.close()
with Reader(
output_file_content(temp_dir,
'*-file_info.jsonl').splitlines()) as jsonl:
record = jsonl.read()
assert '@timestamp' in record
assert record['labels']['artifact'] == "TestArtifact"
assert record['file']['path'].endswith('MSVCR71.dll')
assert record['file']['size'] == 348160
assert record['file']['mime_type'] == "application/x-msdownload"
assert record['file']['hash'][
'md5'] == "86f1895ae8c5e8b17d99ece768a70732"
assert record['file']['hash'][
'sha1'] == "d5502a1d00787d68f548ddeebbde1eca5e2b38ca"
assert record['file']['hash'][
'sha256'] == "8094af5ee310714caebccaeee7769ffb08048503ba478b879edfef5f1a24fefe"
assert record['file']['pe']['company'] == "Microsoft Corporation"
assert record['file']['pe'][
'description'] == "Microsoft® C Runtime Library"
assert record['file']['pe']['file_version'] == "7.10.3052.4"
assert record['file']['pe']['original_file_name'] == "MSVCR71.DLL"
assert record['file']['pe'][
'product'] == "Microsoft® Visual Studio .NET"
assert record['file']['pe'][
'imphash'] == "7acc8c379c768a1ecd81ec502ff5f33e"
assert record['file']['pe']['compilation'] == "2003-02-21T12:42:20"
def test_collect_wmi(temp_dir):
output = Outputs(temp_dir, None, False)
output.add_collected_wmi('TestArtifact', 'query', 'output')
output.close()
wmi = json.loads(output_file_content(temp_dir, '*-wmi.json'))
assert wmi == {'TestArtifact': {'query': 'output'}}
def test_collect_command(temp_dir):
output = Outputs(temp_dir, None, False)
output.add_collected_command('TestArtifact', 'command', b'output')
output.close()
commands = json.loads(output_file_content(temp_dir, '*-commands.json'))
assert commands == {'TestArtifact': {'command': 'output'}}
def test_collect_file(temp_dir, test_file):
output = Outputs(temp_dir, None, False)
output.add_collected_file('TestArtifact',
OSFileSystem('/').get_fullpath(test_file))
output.close()
zip_content = io.BytesIO(output_file_content(temp_dir, '*-files.zip'))
zipfile = ZipFile(zip_content)
zipped_file = zipfile.namelist()[0]
assert zipped_file.endswith('test_file.txt')
def test_collect_registry(temp_dir):
output = Outputs(temp_dir, None, False)
output.add_collected_registry_value('TestArtifact', 'key', 'name', 'value',
'type')
output.close()
registry = json.loads(output_file_content(temp_dir, '*-registry.json'))
assert registry == {
'TestArtifact': {
'key': {
'name': {
'value': 'value',
'type': 'type'
}
}
}
}
def test_collect_file_info(temp_dir, test_file):
output = Outputs(temp_dir, None, False)
output.add_collected_file_info('TestArtifact',
OSFileSystem('/').get_fullpath(test_file))
output.close()
with Reader(
output_file_content(temp_dir,
'*-file_info.jsonl').splitlines()) as jsonl:
record = jsonl.read()
assert '@timestamp' in record
assert record['file']['path'].endswith('test_file.txt')
assert record['file']['size'] == 14
assert record['file']['mime_type'] == "application/x-msdownload"
assert record['file']['hash'][
'md5'] == "10dbf3e392abcc57f8fae061c7c0aeec"
assert record['file']['hash'][
'sha1'] == "7ef0fe6c3855fbac1884e95622d9e45ce1d4ae9b"
assert record['file']['hash'][
'sha256'] == "cfb91ddbf08c52ff294fdf1657081a98c090d270dbb412a91ace815b3df947b6"