def test_getfilerep(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
# Notepad.exe reputations
reputations_dict = \
tie_client.get_file_reputation(FILE_NOTEPAD_EXE_HASH_DICT)
self.assertEqual(
reputations_dict[FileProvider.GTI][
ReputationProp.TRUST_LEVEL], TrustLevel.KNOWN_TRUSTED)
self.assertEqual(
reputations_dict[FileProvider.ENTERPRISE][
ReputationProp.TRUST_LEVEL], TrustLevel.NOT_SET)
# EICAR reputations
reputations_dict = \
tie_client.get_file_reputation(FILE_EICAR_HASH_DICT)
self.assertEqual(
reputations_dict[FileProvider.GTI]
[ReputationProp.TRUST_LEVEL], TrustLevel.KNOWN_MALICIOUS)
self.assertEqual(
reputations_dict[FileProvider.ENTERPRISE][
ReputationProp.TRUST_LEVEL], TrustLevel.NOT_SET)
dxl_client.disconnect()
def test_setcertrep(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
# Set cert1 reputations
# Set the Enterprise reputation for cert1 to Most Likely Trusted
tie_client.set_certificate_reputation(
trust_level=TrustLevel.MOST_LIKELY_TRUSTED,
sha1=CERT_CERT1_SHA1,
public_key_sha1=CERT_CERT1_PUBLIC_KEY_SHA1,
comment=SET_REP_COMMENT)
# Get cert1 reputations
reputations_dict = \
tie_client.get_certificate_reputation(
CERT_CERT1_SHA1,
CERT_CERT1_PUBLIC_KEY_SHA1
)
self.assertEqual(
reputations_dict[CertProvider.ENTERPRISE][
ReputationProp.TRUST_LEVEL],
TrustLevel.MOST_LIKELY_TRUSTED)
dxl_client.disconnect()
def test_setfilerep(self):
file_notepad_exe_filename = "notepad.exe"
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
# Set Notepad.exe reputations
# Set the Enterprise reputation for notepad.exe to Most Likely Trusted
tie_client.set_file_reputation(
TrustLevel.MOST_LIKELY_TRUSTED,
FILE_NOTEPAD_EXE_HASH_DICT,
filename=file_notepad_exe_filename,
comment=SET_REP_COMMENT)
# Get Notepad.exe reputations
reputations_dict = \
tie_client.get_file_reputation(FILE_NOTEPAD_EXE_HASH_DICT)
self.assertEqual(
reputations_dict[FileProvider.ENTERPRISE][
ReputationProp.TRUST_LEVEL],
TrustLevel.MOST_LIKELY_TRUSTED)
dxl_client.disconnect()
def test_basicgetrep_example(self):
# Modify sample file to include necessary sample data
sample_filename = self.BASIC_FOLDER + "/basic_get_reputation_example.py"
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
dxl_client.connect()
with MockTieServer(dxl_client):
mock_print = BaseClientTest.run_sample(sample_filename)
output_string = ""
for print_call in mock_print.mock_calls:
output_string += str(print_call[1]) # Gets text argument for each print call
self.assertTrue(
StringContains('"trustLevel": 99')
)
self.assertTrue(
StringContains('"createDate": 1451502875')
)
self.assertTrue(
StringContains('"providerId": 1')
)
self.assertTrue(
StringContains('"2139285": "72339069014638857"')
)
mock_print.assert_any_call(
StringDoesNotContain("Error")
)
dxl_client.disconnect()
def test_advancedgetrep_example(self):
# Modify sample file to include necessary sample data
sample_filename = self.ADVANCED_FOLDER + "/advanced_get_reputation_example.py"
temp_sample_file = TempSampleFile(sample_filename)
target_line = "FILE_MD5 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[HashType.MD5] + "\"\n"
temp_sample_file.write_file_line(
target=target_line,
replacement=replacement_line
)
target_line = "FILE_SHA1 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[HashType.SHA1] + "\"\n"
temp_sample_file.write_file_line(
target=target_line,
replacement=replacement_line
)
target_line = "FILE_SHA256 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[HashType.SHA256] + "\"\n"
temp_sample_file.write_file_line(
target=target_line,
replacement=replacement_line
)
target_line = "CERTIFICATE_BODY_SHA1 = "
replacement_line = target_line + "\"" + CERT_CERT1_SHA1 + "\"\n"
temp_sample_file.write_file_line(
target=target_line,
replacement=replacement_line
)
target_line = "CERTIFICATE_PUBLIC_KEY_SHA1 = "
replacement_line = target_line + "\"" + CERT_CERT1_PUBLIC_KEY_SHA1 + "\"\n"
temp_sample_file.write_file_line(
target=target_line,
replacement=replacement_line
)
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
dxl_client.connect()
with MockTieServer(dxl_client):
mock_print = BaseClientTest.run_sample(temp_sample_file.temp_file.name)
mock_print.assert_any_call(
StringContains(GtiAttrib.ORIGINAL_RESPONSE)
)
mock_print.assert_any_call(
StringContains(FileEnterpriseAttrib.FIRST_CONTACT)
)
mock_print.assert_any_call(
StringDoesNotContain("Error")
)
dxl_client.disconnect()
def test_basicfirstref_example(self):
# Modify sample file to include necessary sample data
sample_filename = self.BASIC_FOLDER + "/basic_first_ref_example.py"
temp_sample_file = TempSampleFile(sample_filename)
target_line = "FILE_MD5 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[
HashType.MD5] + "\"\n"
temp_sample_file.write_file_line(target_line, replacement_line)
target_line = "FILE_SHA1 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[
HashType.SHA1] + "\"\n"
temp_sample_file.write_file_line(target_line, replacement_line)
target_line = "FILE_SHA256 = "
replacement_line = target_line + "\"" + FILE_NOTEPAD_EXE_HASH_DICT[
HashType.SHA256] + "\"\n"
temp_sample_file.write_file_line(target_line, replacement_line)
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
dxl_client.connect()
with MockTieServer(dxl_client):
mock_print = BaseClientTest.run_sample(
temp_sample_file.temp_file.name)
for guid in FIRST_REF_AGENT_GUIDS:
mock_print.assert_any_call(StringContains(guid))
mock_print.assert_any_call(StringDoesNotContain("Error"))
dxl_client.disconnect()
def test_getfilerep_invalid(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
self.assertRaisesRegex(
Exception, r"Error: Could not find reputation \(0\)",
tie_client.get_file_reputation, FILE_INVALID_HASH_DICT)
dxl_client.disconnect()
def test_getcertfirstrefs_invalid(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
self.assertRaisesRegex(
Exception, r"Error: Could not find reputation \(0\)",
tie_client.get_certificate_first_references,
CERT_INVALID_SHA1, CERT_INVALID_SHA1)
dxl_client.disconnect()
def test_getfilefirstrefs(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
# Notepad.exe reputations
systems_list = \
tie_client.get_file_first_references(FILE_NOTEPAD_EXE_HASH_DICT)
for system in systems_list:
self.assertIn(system[FirstRefProp.SYSTEM_GUID],
FIRST_REF_AGENT_GUIDS)
dxl_client.disconnect()
def test_getcertfirstrefs(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
systems_list = \
tie_client.get_certificate_first_references(
CERT_CERT1_SHA1,
CERT_CERT1_PUBLIC_KEY_SHA1
)
for system in systems_list:
self.assertIn(system[FirstRefProp.SYSTEM_GUID],
FIRST_REF_AGENT_GUIDS)
dxl_client.disconnect()
def test_basicsetrep_example(self):
# Modify sample file to include necessary sample data
sample_filename = self.BASIC_FOLDER + "/basic_set_reputation_example.py"
temp_sample_file = TempSampleFile(sample_filename)
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
dxl_client.connect()
with MockTieServer(dxl_client):
mock_print = BaseClientTest.run_sample(
temp_sample_file.temp_file.name)
mock_print.assert_any_call(StringContains("Succeeded"))
mock_print.assert_any_call(StringDoesNotContain("Error"))
dxl_client.disconnect()
def test_basicsetexternalrep_example_succeed(self):
# Modify sample file to include necessary sample data
sample_filename = self.BASIC_FOLDER + "/basic_set_external_file_reputation.py"
temp_sample_file = TempSampleFile(sample_filename)
target_line = "fileMD5 = "
replacement_line = target_line + "\"" + FILE_UNKNOWN_HASH_DICT[HashType.MD5] + "\"\n"
temp_sample_file.write_file_line(
target_line,
replacement_line
)
target_line = "fileSHA1 = "
replacement_line = target_line + "\"" + FILE_UNKNOWN_HASH_DICT[HashType.SHA1] + "\"\n"
temp_sample_file.write_file_line(
target_line,
replacement_line
)
target_line = "fileSHA256 = "
replacement_line = target_line + "\"" + FILE_UNKNOWN_HASH_DICT[HashType.SHA256] + "\"\n"
temp_sample_file.write_file_line(
target_line,
replacement_line
)
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
dxl_client.connect()
with MockTieServer(dxl_client):
mock_print = BaseClientTest.run_sample(temp_sample_file.temp_file.name)
mock_print.assert_any_call(
StringContains("Event Sent")
)
mock_print.assert_any_call(
StringDoesNotContain("Error")
)
dxl_client.disconnect()
def test_getcertrep(self):
with self.create_client(max_retries=0) as dxl_client:
# Set up client, and register mock service
tie_client = TieClient(dxl_client)
dxl_client.connect()
with MockTieServer(dxl_client):
# cert1 reputations
reputations_dict = \
tie_client.get_certificate_reputation(
CERT_CERT1_SHA1,
CERT_CERT1_PUBLIC_KEY_SHA1
)
self.assertEqual(
reputations_dict[CertProvider.GTI][
ReputationProp.TRUST_LEVEL], TrustLevel.KNOWN_TRUSTED)
self.assertEqual(
reputations_dict[CertProvider.ENTERPRISE][
ReputationProp.TRUST_LEVEL], TrustLevel.NOT_SET)
dxl_client.disconnect()
