Skip to content

Commit 4137037

Browse files
calebmadrigalcalebmadrigal
committed
Factored out config management, unit tests
1 parent 3d8e938 commit 4137037

5 files changed

Lines changed: 243 additions & 205 deletions

File tree

‎tests/test_config_management.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# pylint: disable=C0111, C0413, C0103, E0401
2+
import unittest
3+
import trackerjacker.config_management as cm
4+
5+
6+
class TestConfigManagement(unittest.TestCase):
7+
def test_parse_watch_list(self):
8+
test1 = 'aa:bb:cc:dd:ee:ff'
9+
parsed = cm.parse_watch_list(test1)
10+
self.assertEqual(parsed, {'aa:bb:cc:dd:ee:ff': 1})
11+
12+
test2 = 'aa:bb:cc:dd:ee:ff, 11:22:33:44:55:66'
13+
parsed = cm.parse_watch_list(test2)
14+
self.assertEqual(parsed, {'aa:bb:cc:dd:ee:ff': 1, '11:22:33:44:55:66': 1})
15+
16+
test3 = 'aa:bb:cc:dd:ee:ff=1000, 11:22:33:44:55:66'
17+
parsed = cm.parse_watch_list(test3)
18+
self.assertEqual(parsed, {'aa:bb:cc:dd:ee:ff': 1000, '11:22:33:44:55:66': 1})
19+
20+
def test_default_config(self):
21+
# Just making sure I understand how parse_args works
22+
cmd_line_args = cm.get_arg_parser().parse_args([])
23+
self.assertEqual(cmd_line_args.do_map, False)
24+
cmd_line_args = cm.get_arg_parser().parse_args(['--map'])
25+
self.assertEqual(cmd_line_args.do_map, True)
26+
27+
cmd_line_args = cm.get_arg_parser().parse_args([])
28+
config = cm.build_config(cmd_line_args)
29+
self.assertEqual(config, cm.DEFAULT_CONFIG)
30+
self.assertEqual(config['map_file'], 'wifi_map.yaml')
31+
32+
def test_override_config(self):
33+
cmd_line_args = cm.get_arg_parser().parse_args(['--map-file', 'my_network.yaml'])
34+
config = cm.build_config(cmd_line_args)
35+
self.assertEqual(config['map_file'], 'my_network.yaml')
36+
37+
def test_config_macs_to_watch(self):
38+
cmd_line_args = cm.get_arg_parser().parse_args(['--track', '-m', '7C:70:BC:78:70:21'])
39+
config = cm.build_config(cmd_line_args)
40+
self.assertEqual(config['devices_to_watch'], {'7C:70:BC:78:70:21': 1})
41+
42+
cmd_line_args = cm.get_arg_parser().parse_args(['--track', '-m', '7C:70:BC:78:70:21=100'])
43+
config = cm.build_config(cmd_line_args)
44+
self.assertEqual(config['devices_to_watch'], {'7C:70:BC:78:70:21': 100})
45+
46+
cmd_line_args = cm.get_arg_parser().parse_args(['--track', '-m', '7C:70:BC:78:70:21=100,aa:bb:cc:dd:ee:ff'])
47+
config = cm.build_config(cmd_line_args)
48+
self.assertEqual(config['devices_to_watch'], {'7C:70:BC:78:70:21': 100, 'aa:bb:cc:dd:ee:ff': 1})
49+
50+
51+
if __name__ == '__main__':
52+
unittest.main()

‎tests/test_main.py‎

Lines changed: 0 additions & 22 deletions
This file was deleted.

‎trackerjacker/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@
77
__author__ = "Caleb Madrigal"
88
__email__ = "[email protected]"
99
__license__ = "MIT"
10-
__version__ = "1.0.3"
10+
__version__ = "1.0.4"

‎trackerjacker/__main__.py‎

Lines changed: 7 additions & 182 deletions
Original file line numberDiff line numberDiff line change
@@ -8,34 +8,15 @@
88
import errno
99
import pprint
1010
import logging
11-
import argparse
1211

12+
from . import config_management
1313
from . import device_management
1414
from . import dot11_frame
1515
from . import dot11_mapper
1616
from . import dot11_tracker
1717
from . import ieee_mac_vendor_db
1818
from .common import TJException
1919

20-
# Default config
21-
DEFAULT_CONFIG = {'log_path': None,
22-
'log_level': 'INFO',
23-
'iface': None,
24-
'devices_to_watch': [],
25-
'aps_to_watch': [],
26-
'threshold_window': 10,
27-
'do_map': True,
28-
'do_track': False,
29-
'map_file': 'wifi_map.yaml',
30-
'map_save_interval': 10,
31-
'alert_cooldown': 30,
32-
'alert_command': None,
33-
'channels_to_monitor': None,
34-
'channel_switch_scheme': 'round_robin',
35-
'time_per_channel': 2,
36-
'display_matching_packets': False,
37-
'display_all_packets': False}
38-
3920
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
4021
try:
4122
import scapy.all as scapy
@@ -204,56 +185,6 @@ def stop(self):
204185
self.dot11_tracker.stop()
205186

206187

207-
def parse_command_line_args():
208-
parser = argparse.ArgumentParser()
209-
# Modes
210-
parser.add_argument('--map', action='store_true', dest='do_map',
211-
help='Map mode - output map to wifi_map.yaml')
212-
parser.add_argument('--track', action='store_true', dest='do_track',
213-
help='Track mode')
214-
parser.add_argument('--monitor-mode-on', action='store_true', dest='do_enable_monitor_mode',
215-
help='Enables monitor mode on the specified interface and exit')
216-
parser.add_argument('--monitor-mode-off', action='store_true', dest='do_disable_monitor_mode',
217-
help='Disables monitor mode on the specified interface and exit')
218-
parser.add_argument('--set-channel', metavar='CHANNEL', dest='set_channel', nargs=1,
219-
help='Set the specified wireless interface to the specified channel and exit')
220-
parser.add_argument('--mac-lookup', type=str, dest='mac_lookup',
221-
help='Lookup the vendor of the specified MAC address and exit')
222-
parser.add_argument('--print-default-config', action='store_true', dest='print_default_config',
223-
help='Print boilerplate config file and exit')
224-
225-
# Normal switches
226-
parser.add_argument('-i', '--interface', type=str, dest='iface',
227-
help='Network interface to use; if empty, try to find monitor inferface')
228-
parser.add_argument('-m', '--macs', type=str, dest='devices_to_watch',
229-
help='MAC(s) to track; comma separated for multiple')
230-
parser.add_argument('-a', '--access-points', type=str, dest='aps_to_watch',
231-
help='Access point(s) to track - specified by BSSID; comma separated for multiple')
232-
parser.add_argument('--channels-to-monitor', type=str, dest='channels_to_monitor',
233-
help='Channels to monitor; comma separated for multiple')
234-
parser.add_argument('-w', '--time-window', type=int, dest='threshold_window',
235-
help='Time window (in seconds) which alert threshold is applied to')
236-
parser.add_argument('--map-save-interval', type=float, dest='map_save_interval',
237-
help='Number of seconds between saving the wifi map to disk')
238-
parser.add_argument('--eval_interval', type=float, dest='eval_interval',
239-
help='Number of seconds between looking for tracked devices')
240-
parser.add_argument('--power', action='store_true', dest='threshold_is_power',
241-
help='If specified, all tracking thresholds are taken to represent RSSI power levels')
242-
parser.add_argument('--alert-command', type=str, dest='alert_command',
243-
help='Command to execute upon alert')
244-
parser.add_argument('--display-all-packets', action='store_true', dest='display_all_packets',
245-
help='If true, displays all packets matching filters')
246-
parser.add_argument('--log-path', type=str, dest='log_path', default=None,
247-
help='Log path; default is stdout')
248-
parser.add_argument('--log-level', type=str, dest='log_level', default='INFO',
249-
help='Log level; Options: DEBUG, INFO, WARNING, ERROR, CRITICAL')
250-
parser.add_argument('-c', '--config', type=str, dest='config',
251-
help='Path to config json file; For example config file, use --print-default-config')
252-
253-
# vars converts from namespace to dict
254-
return parser.parse_args()
255-
256-
257188
def do_simple_tasks_if_specified(args):
258189
if args.do_enable_monitor_mode:
259190
if not args.iface:
@@ -275,7 +206,7 @@ def do_simple_tasks_if_specified(args):
275206
print('Vendor for {} not found'.format(args.mac_lookup), file=sys.stderr)
276207
sys.exit(0)
277208
elif args.print_default_config:
278-
print(json.dumps(DEFAULT_CONFIG, indent=4, sort_keys=True))
209+
print(json.dumps(config_management.DEFAULT_CONFIG, indent=4, sort_keys=True))
279210
sys.exit(0)
280211
elif args.set_channel:
281212
if not args.iface:
@@ -286,121 +217,12 @@ def do_simple_tasks_if_specified(args):
286217
sys.exit(0)
287218

288219

289-
def build_config(args):
290-
config = DEFAULT_CONFIG
291-
292-
macs_from_config = []
293-
aps_from_config = []
294-
295-
if args.config:
296-
try:
297-
with open(args.config, 'r') as f:
298-
config_from_file = json.loads(f.read())
299-
300-
# If there are any keys defined in the config file not allowed, error out
301-
invalid_keys = set(config_from_file.keys()) - set(config.keys())
302-
if invalid_keys:
303-
raise TJException('Invalid keys found in config file: {}'.format(invalid_keys))
304-
305-
macs_from_config = config_from_file.pop('devices_to_watch', {})
306-
aps_from_config = config_from_file.pop('aps_to_watch', {})
307-
308-
config.update(config_from_file)
309-
print('Loaded configuration from {}'.format(args.config))
310-
311-
except (IOError, OSError, json.decoder.JSONDecodeError) as e:
312-
raise TJException('Error loading config file ({}): {}'.format(args.config, e))
313-
314-
macs_from_args = {}
315-
aps_from_args = {}
316-
317-
# Converts from cli param format like: "aa:bb:cc:dd:ee:ff,11:22:33:44:55:66' to a map like
318-
# {'aa:bb:cc:dd:ee:ff': 1, '11:22:33:44:55:66': 1}
319-
if args.devices_to_watch:
320-
macs_from_args = parse_watch_list(args.devices_to_watch)
321-
322-
# Converts from cli param format like "my_ssid1=5000,bssid2=1337" to a map like:
323-
# {'my_ssid1': 5000, 'bssid2': 1337}
324-
if args.aps_to_watch:
325-
aps_from_args = parse_watch_list(args.aps_to_watch)
326-
327-
non_config_args = ['config', 'devices_to_watch', 'aps_to_watch', 'do_enable_monitor_mode',
328-
'do_disable_monitor_mode', 'set_channel', 'print_default_config', 'mac_lookup']
329-
330-
config_from_args = vars(args)
331-
config_from_args = {k: v for k, v in config_from_args.items()
332-
if v is not None and k not in non_config_args}
333-
334-
# Config from args trumps everything
335-
config.update(config_from_args)
336-
337-
# Only allow track or map mode at once
338-
if config['do_track']:
339-
config['do_map'] = False
340-
if not config['do_track']:
341-
config['do_map'] = True
342-
343-
config['devices_to_watch'] = dict(macs_from_config, **macs_from_args)
344-
config['aps_to_watch'] = dict(aps_from_config, **aps_from_args)
345-
346-
if args.channels_to_monitor:
347-
channels_to_monitor = args.channels_to_monitor.split(',')
348-
config['channels_to_monitor'] = channels_to_monitor
349-
350-
if config['log_level'] == 'DEBUG':
351-
print('Config:')
352-
pprint.pprint(config)
353-
354-
return config
355-
356-
357-
def parse_watch_list(watch_str):
358-
""" Parse string to represent devices to watch config
359-
360-
Valid examples:
361-
* aa:bb:cc:dd:ee:ff
362-
- Threshold of 1 for the given MAC address
363-
* aa:bb:cc:dd:ee:ff,11:22:33:44:55:66
364-
- This means look for any traffic from either address
365-
* aa:bb:cc:dd:ee:ff=1337, 11:22:33:44:55:66=1000
366-
- This means look for 1337 bytes for the first address, and 1000 for the second
367-
* my_ssid, 11:22:33:44:55:66=1000
368-
- This means look for 1 byte from my_ssid or 1000 for the second
369-
370-
Returns dict in this format:
371-
{'aa:bb:cc:dd:ee:ff': threshold1, '11:22:33:44:55:66': threshold2}
372-
"""
373-
374-
watch_list = [i.strip() for i in watch_str.split(',')]
375-
watch_dict = {}
376-
377-
for watch_part in watch_list:
378-
if '=' in watch_part:
379-
# dev_id is a MAC, BSSID, or SSID
380-
dev_id, threshold = [i.strip() for i in watch_part.split('=')]
381-
try:
382-
threshold = int(threshold)
383-
except ValueError:
384-
# Can't parse with "dev_id=threshold" formula, so assume '=' sign was part of ssid
385-
dev_id = watch_part
386-
threshold = 1
387-
388-
watch_part.split('=')
389-
else:
390-
# Can't parse with "dev_id=threshold" formula, so assume...
391-
dev_id = watch_part
392-
threshold = 1
393-
394-
watch_dict[dev_id] = threshold
395-
return watch_dict
396-
397-
398220
def main():
399221
if not os.getuid() == 0:
400222
print('trackerjacker requires r00t!', file=sys.stderr)
401223
sys.exit(errno.EPERM)
402224

403-
argparse_args = parse_command_line_args()
225+
argparse_args = config_management.get_arg_parser().parse_args()
404226

405227
# Some command-line args specify to just perform a simple task and then exit
406228
try:
@@ -409,7 +231,10 @@ def main():
409231
print('Error: {}'.format(e), file=sys.stderr)
410232
sys.exit(1)
411233

412-
config = build_config(argparse_args)
234+
config = config_management.build_config(argparse_args)
235+
if config['log_level'] == 'DEBUG':
236+
print('Config:')
237+
pprint.pprint(config)
413238

414239
# Setup logger
415240
logger = make_logger(config.pop('log_path'), config.pop('log_level'))

0 commit comments

Comments
 (0)