Skip to content

Commit f05128c

Browse files
kimocoderkimocoder
committed
session: add session/resume/clean support
1 parent f7fb247 commit f05128c

20 files changed

Lines changed: 5417 additions & 33 deletions

‎tests/test_config_restoration.py‎

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,272 @@
1+
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
3+
4+
"""
5+
Tests for configuration restoration from session.
6+
"""
7+
8+
import unittest
9+
import tempfile
10+
import shutil
11+
import os
12+
from unittest.mock import Mock, patch, MagicMock
13+
from wifite.util.session import SessionManager, SessionState, TargetState
14+
15+
16+
class TestConfigurationRestoration(unittest.TestCase):
17+
"""Test configuration restoration from session."""
18+
19+
def setUp(self):
20+
"""Set up test fixtures."""
21+
self.temp_dir = tempfile.mkdtemp()
22+
self.session_mgr = SessionManager(session_dir=self.temp_dir)
23+
24+
# Create a mock Configuration object
25+
self.config = Mock()
26+
self.config.interface = 'wlan0mon'
27+
self.config.wordlist = '/usr/share/wordlists/rockyou.txt'
28+
self.config.wpa_attack_timeout = 500
29+
self.config.wps_pixie = True
30+
self.config.wps_pin = True
31+
self.config.dont_use_pmkid = False
32+
self.config.wps_only = False
33+
self.config.use_pmkid_only = False
34+
self.config.infinite_mode = False
35+
self.config.attack_max = 0
36+
self.config.use_tui = True
37+
self.config.verbose = 0
38+
39+
def tearDown(self):
40+
"""Clean up test fixtures."""
41+
shutil.rmtree(self.temp_dir)
42+
43+
def test_restore_basic_configuration(self):
44+
"""Test restoring basic configuration parameters."""
45+
# Set config to None values to avoid conflicts
46+
self.config.interface = None
47+
self.config.wordlist = None
48+
self.config.use_tui = None
49+
50+
# Create a session with specific configuration
51+
session = SessionState(
52+
session_id='test_session',
53+
created_at=1234567890.0,
54+
updated_at=1234567890.0,
55+
config={
56+
'interface': 'wlan1mon',
57+
'wordlist': '/custom/wordlist.txt',
58+
'wpa_attack_timeout': 600,
59+
'wps_pixie': False,
60+
'wps_pin': False,
61+
'use_pmkid': True,
62+
'wps_only': False,
63+
'use_pmkid_only': True,
64+
'infinite_mode': True,
65+
'attack_max': 5,
66+
'use_tui': False,
67+
'verbose': 2
68+
},
69+
targets=[
70+
TargetState(
71+
bssid='AA:BB:CC:DD:EE:FF',
72+
essid='TestNetwork',
73+
channel=6,
74+
encryption='WPA2',
75+
power=50,
76+
wps=False
77+
)
78+
]
79+
)
80+
81+
# Mock subprocess to simulate interface check
82+
with patch('subprocess.run') as mock_run:
83+
mock_run.return_value = Mock(
84+
stdout='Interface wlan1mon\n',
85+
returncode=0
86+
)
87+
88+
# Restore configuration
89+
result = self.session_mgr.restore_configuration(session, self.config)
90+
91+
# Verify configuration was restored
92+
self.assertEqual(self.config.interface, 'wlan1mon')
93+
self.assertEqual(self.config.wordlist, '/custom/wordlist.txt')
94+
self.assertEqual(self.config.wpa_attack_timeout, 600)
95+
self.assertFalse(self.config.wps_pixie)
96+
self.assertFalse(self.config.wps_pin)
97+
self.assertFalse(self.config.dont_use_pmkid) # use_pmkid=True means dont_use_pmkid=False
98+
self.assertTrue(self.config.use_pmkid_only)
99+
self.assertTrue(self.config.infinite_mode)
100+
self.assertEqual(self.config.attack_max, 5)
101+
self.assertFalse(self.config.use_tui)
102+
self.assertEqual(self.config.verbose, 2)
103+
104+
# Should have no warnings (interface is available)
105+
self.assertEqual(len(result['warnings']), 0)
106+
# May have conflicts if default values differ, but that's okay
107+
self.assertFalse(result['interface_changed'])
108+
109+
def test_interface_not_available(self):
110+
"""Test handling when saved interface is not available."""
111+
session = SessionState(
112+
session_id='test_session',
113+
created_at=1234567890.0,
114+
updated_at=1234567890.0,
115+
config={
116+
'interface': 'wlan5mon', # Non-existent interface
117+
'wordlist': '/usr/share/wordlists/rockyou.txt'
118+
},
119+
targets=[
120+
TargetState(
121+
bssid='AA:BB:CC:DD:EE:FF',
122+
essid='TestNetwork',
123+
channel=6,
124+
encryption='WPA2',
125+
power=50,
126+
wps=False
127+
)
128+
]
129+
)
130+
131+
# Mock subprocess to simulate interface not found
132+
with patch('subprocess.run') as mock_run:
133+
mock_run.return_value = Mock(
134+
stdout='Interface wlan0mon\nInterface wlan1mon\n',
135+
returncode=0
136+
)
137+
138+
# Restore configuration
139+
result = self.session_mgr.restore_configuration(session, self.config)
140+
141+
# Should have warnings about interface
142+
self.assertGreater(len(result['warnings']), 0)
143+
self.assertTrue(any('wlan5mon' in w for w in result['warnings']))
144+
self.assertTrue(result['interface_changed'])
145+
146+
def test_conflicting_command_line_flags(self):
147+
"""Test detection of conflicting command-line flags."""
148+
# Set different values in config (simulating command-line flags)
149+
self.config.wordlist = '/different/wordlist.txt'
150+
self.config.wpa_attack_timeout = 300
151+
self.config.use_tui = False
152+
153+
session = SessionState(
154+
session_id='test_session',
155+
created_at=1234567890.0,
156+
updated_at=1234567890.0,
157+
config={
158+
'interface': 'wlan0mon',
159+
'wordlist': '/original/wordlist.txt',
160+
'wpa_attack_timeout': 600,
161+
'use_tui': True
162+
},
163+
targets=[
164+
TargetState(
165+
bssid='AA:BB:CC:DD:EE:FF',
166+
essid='TestNetwork',
167+
channel=6,
168+
encryption='WPA2',
169+
power=50,
170+
wps=False
171+
)
172+
]
173+
)
174+
175+
# Mock subprocess
176+
with patch('subprocess.run') as mock_run:
177+
mock_run.return_value = Mock(
178+
stdout='Interface wlan0mon\n',
179+
returncode=0
180+
)
181+
182+
# Restore configuration
183+
result = self.session_mgr.restore_configuration(session, self.config)
184+
185+
# Should have conflicts detected
186+
self.assertGreater(len(result['conflicts']), 0)
187+
188+
# Configuration should be overridden with session values
189+
self.assertEqual(self.config.wordlist, '/original/wordlist.txt')
190+
self.assertEqual(self.config.wpa_attack_timeout, 600)
191+
self.assertTrue(self.config.use_tui)
192+
193+
def test_interface_check_failure(self):
194+
"""Test handling when interface check fails."""
195+
session = SessionState(
196+
session_id='test_session',
197+
created_at=1234567890.0,
198+
updated_at=1234567890.0,
199+
config={
200+
'interface': 'wlan0mon',
201+
'wordlist': '/usr/share/wordlists/rockyou.txt'
202+
},
203+
targets=[
204+
TargetState(
205+
bssid='AA:BB:CC:DD:EE:FF',
206+
essid='TestNetwork',
207+
channel=6,
208+
encryption='WPA2',
209+
power=50,
210+
wps=False
211+
)
212+
]
213+
)
214+
215+
# Mock subprocess to raise exception
216+
with patch('subprocess.run') as mock_run:
217+
mock_run.side_effect = FileNotFoundError('iw command not found')
218+
219+
# Restore configuration
220+
result = self.session_mgr.restore_configuration(session, self.config)
221+
222+
# Should have warning about interface check failure
223+
self.assertGreater(len(result['warnings']), 0)
224+
self.assertTrue(any('verify interface' in w for w in result['warnings']))
225+
self.assertTrue(result['interface_changed'])
226+
227+
def test_restore_with_missing_config_values(self):
228+
"""Test restoration when some config values are missing."""
229+
session = SessionState(
230+
session_id='test_session',
231+
created_at=1234567890.0,
232+
updated_at=1234567890.0,
233+
config={
234+
'interface': 'wlan0mon',
235+
# Missing wordlist, timeouts, etc.
236+
},
237+
targets=[
238+
TargetState(
239+
bssid='AA:BB:CC:DD:EE:FF',
240+
essid='TestNetwork',
241+
channel=6,
242+
encryption='WPA2',
243+
power=50,
244+
wps=False
245+
)
246+
]
247+
)
248+
249+
# Store original values
250+
original_wordlist = self.config.wordlist
251+
original_timeout = self.config.wpa_attack_timeout
252+
253+
# Mock subprocess
254+
with patch('subprocess.run') as mock_run:
255+
mock_run.return_value = Mock(
256+
stdout='Interface wlan0mon\n',
257+
returncode=0
258+
)
259+
260+
# Restore configuration
261+
result = self.session_mgr.restore_configuration(session, self.config)
262+
263+
# Original values should be preserved when not in session
264+
self.assertEqual(self.config.wordlist, original_wordlist)
265+
self.assertEqual(self.config.wpa_attack_timeout, original_timeout)
266+
267+
# Interface should still be restored
268+
self.assertEqual(self.config.interface, 'wlan0mon')
269+
270+
271+
if __name__ == '__main__':
272+
unittest.main()

0 commit comments

Comments
 (0)