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()