def test_model_generation(self):
"""Ensure a model is generated from a summary file and it is
correct"""
with self.mktemp() as output:
summary = os.path.join(DATA_DIR, "traffic-sample-very-short.txt")
command = "%s %s --out %s" % (LEARNER, summary, output)
self.check_run(command)
expected_fn = os.path.join(DATA_DIR, "traffic_learner.expected")
expected = traffic.TrafficModel()
f=open(expected_fn)
expected.load(f)
f.close()
f=open(output)
actual = traffic.TrafficModel()
actual.load(f)
f.close()
actual_ngrams = {k: sorted(v) for k, v in actual.ngrams.items()}
expected_ngrams = {k: sorted(v) for k, v in expected.ngrams.items()}
self.assertEquals(expected_ngrams, actual_ngrams)
actual_details = {k: sorted(v) for k, v in actual.query_details.items()}
expected_details = {k: sorted(v) for k, v in expected.query_details.items()}
self.assertEquals(expected_details, actual_details)
self.assertEquals(expected.cumulative_duration, actual.cumulative_duration)
self.assertEquals(expected.packet_rate, actual.packet_rate)
with open(expected_fn) as f1, open(output) as f2:
expected_json = json.load(f1)
actual_json = json.load(f2)
self.assertEqual(expected_json, actual_json)
def test_parse_ngrams_dns_included(self):
model = traffic.TrafficModel()
f = open(TEST_FILE)
(conversations, interval, duration,
dns_counts) = traffic.ingest_summaries([f], dns_mode='include')
f.close()
model.learn(conversations)
expected_ngrams = {
('-', '-'): ['dns:0', 'dns:0', 'dns:0', 'ldap:3'],
('-', 'dns:0'): ['dns:0', 'dns:0', 'dns:0'],
('-', 'ldap:3'): ['wait:0'],
('cldap:3', 'cldap:3'): ['cldap:3', 'wait:0'],
('cldap:3', 'wait:0'): ['rpc_netlogon:29'],
('dns:0', 'dns:0'): ['dns:0', 'dns:0', 'dns:0', 'wait:0'],
('dns:0', 'wait:0'): ['cldap:3'],
('kerberos:', 'ldap:3'): ['-'],
('ldap:3', 'wait:0'): ['ldap:2'],
('rpc_netlogon:29', 'kerberos:'): ['ldap:3'],
('wait:0', 'cldap:3'): ['cldap:3'],
('wait:0', 'rpc_netlogon:29'): ['kerberos:']
}
expected_query_details = {
'cldap:3': [('', '', '', 'Netlogon', '', '', ''),
('', '', '', 'Netlogon', '', '', ''),
('', '', '', 'Netlogon', '', '', '')],
'dns:0': [(), (), (), (), (), (), (), (), ()],
'kerberos:': [('', )],
'ldap:2': [('', '', '', '', '', '', '')],
'ldap:3':
[('', '', '', 'subschemaSubentry,dsServiceName,namingContexts,'
'defaultNamingContext,schemaNamingContext,'
'configurationNamingContext,rootDomainNamingContext,'
'supportedControl,supportedLDAPVersion,'
'supportedLDAPPolicies,supportedSASLMechanisms,'
'dnsHostName,ldapServiceName,serverName,'
'supportedCapabilities', '', '', ''),
('2', 'DC,DC', '', 'cn', '', '', '')],
'rpc_netlogon:29': [()]
}
self.maxDiff = 5000
ngrams = {k: sorted(v) for k, v in model.ngrams.items()}
details = {k: sorted(v) for k, v in model.query_details.items()}
self.assertEqual(expected_ngrams, ngrams)
self.assertEqual(expected_query_details, details)
# We use a stringIO instead of a temporary file
f = StringIO()
model.save(f)
model2 = traffic.TrafficModel()
f.seek(0)
model2.load(f)
ngrams = {k: sorted(v) for k, v in model2.ngrams.items()}
details = {k: sorted(v) for k, v in model2.query_details.items()}
self.assertEqual(expected_ngrams, ngrams)
self.assertEqual(expected_query_details, details)
def setUp(self):
self.model = traffic.TrafficModel()