1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
|
# -*- coding: utf-8 -*-
"""Tests for the source type objects."""
import unittest
from artifacts import errors
from artifacts import source_type
class TestSourceType(source_type.SourceType):
"""Class that implements a test source type."""
TYPE_INDICATOR = u'test'
def __init__(self, test=None):
"""Initializes the source type object.
Args:
test: optional test string. The default is None.
Raises:
FormatError: when test is not set.
"""
if not test:
raise errors.FormatError(u'Missing test value.')
super(TestSourceType, self).__init__()
self.test = test
def CopyToDict(self):
"""Copies the source type to a dictionary.
Returns:
A dictionary containing the source type attributes.
"""
return {u'test': self.test}
class SourceTypeTest(unittest.TestCase):
"""Class to test the artifact source type."""
class ArtifactGroupSourceTypeTest(unittest.TestCase):
"""Class to test the artifact group source type."""
def testInitialize(self):
"""Tests the __init__ function."""
source_type.ArtifactGroupSourceType(names=[u'test'])
class FileSourceTypeTest(unittest.TestCase):
"""Class to test the files source type."""
def testInitialize(self):
"""Tests the __init__ function."""
source_type.FileSourceType(paths=[u'test'])
source_type.FileSourceType(paths=[u'test'], separator=u'\\')
class PathSourceTypeTest(unittest.TestCase):
"""Class to test the paths source type."""
def testInitialize(self):
"""Tests the __init__ function."""
source_type.PathSourceType(paths=[u'test'])
source_type.PathSourceType(paths=[u'test'], separator=u'\\')
class WindowsRegistryKeySourceTypeTest(unittest.TestCase):
"""Class to test the Windows Registry keys source type."""
def testInitialize(self):
"""Tests the __init__ function."""
source_type.WindowsRegistryKeySourceType(keys=[u'HKEY_LOCAL_MACHINE\\test'])
with self.assertRaises(errors.FormatError):
source_type.WindowsRegistryKeySourceType(keys=u'HKEY_LOCAL_MACHINE\\test')
class WindowsRegistryValueSourceTypeTest(unittest.TestCase):
"""Class to test the Windows Registry value source type."""
def testInitialize(self):
"""Tests the __init__ function."""
key_value_pair = {'key': u'HKEY_LOCAL_MACHINE\\test', 'value': u'test'}
source_type.WindowsRegistryValueSourceType(key_value_pairs=[key_value_pair])
key_value_pair = {'bad': u'test', 'value': u'test'}
with self.assertRaises(errors.FormatError):
source_type.WindowsRegistryValueSourceType(
key_value_pairs=[key_value_pair])
with self.assertRaises(errors.FormatError):
source_type.WindowsRegistryValueSourceType(key_value_pairs=key_value_pair)
class WMIQuerySourceTypeTest(unittest.TestCase):
"""Class to test the WMI query source type."""
def testInitialize(self):
"""Tests the __init__ function."""
source_type.WMIQuerySourceType(query=u'test')
class SourceTypeFactoryTest(unittest.TestCase):
"""Class to test the source type factory."""
def testCreateSourceType(self):
"""Tests the source type creation."""
source_type.SourceTypeFactory.RegisterSourceTypes([TestSourceType])
with self.assertRaises(KeyError):
source_type.SourceTypeFactory.RegisterSourceTypes([TestSourceType])
source_object = source_type.SourceTypeFactory.CreateSourceType(
u'test', {u'test': u'test123'})
self.assertIsNotNone(source_object)
self.assertEqual(source_object.test, u'test123')
with self.assertRaises(errors.FormatError):
source_object = source_type.SourceTypeFactory.CreateSourceType(u'test',
{})
with self.assertRaises(errors.FormatError):
source_object = source_type.SourceTypeFactory.CreateSourceType(u'bogus',
{})
source_type.SourceTypeFactory.DeregisterSourceType(TestSourceType)
def testRegisterSourceType(self):
"""Tests the source type registration functions."""
expected_number_of_source_types = len(
source_type.SourceTypeFactory.GetSourceTypes())
source_type.SourceTypeFactory.RegisterSourceType(TestSourceType)
number_of_source_types = len(source_type.SourceTypeFactory.GetSourceTypes())
self.assertEqual(number_of_source_types,
expected_number_of_source_types + 1)
source_type.SourceTypeFactory.DeregisterSourceType(TestSourceType)
number_of_source_types = len(source_type.SourceTypeFactory.GetSourceTypes())
self.assertEqual(number_of_source_types, expected_number_of_source_types)
def testRegisterSourceTypeRaisesWhenAlreadyRegistered(self):
"""Tests the source type registration functions when already registered."""
source_type.SourceTypeFactory.RegisterSourceType(TestSourceType)
with self.assertRaises(KeyError):
source_type.SourceTypeFactory.RegisterSourceType(TestSourceType)
source_type.SourceTypeFactory.DeregisterSourceType(TestSourceType)
if __name__ == '__main__':
unittest.main()
|
