def test_action_space(self):
actual_A_space = Action.generate_action_space(self.ads_space, self.E,
self.test_exploit_cost, self.test_scan_cost)
expected_A_space = []
for m in self.ads_space:
scan = Action(m, self.test_scan_cost, "scan")
expected_A_space.append(scan)
for e in range(self.E):
exploit = Action(m, self.test_exploit_cost, "exploit", e)
expected_A_space.append(exploit)
self.assertEqual(actual_A_space, expected_A_space)
def test_invalid_action(self):
# invalid subnet
exploit = Action((USER + 1, 0), A_COST, "scan", None)
with self.assertRaises(AssertionError):
self.network.perform_action(exploit)
# invalid machine ID
exploit = Action((DMZ, 2), A_COST, "scan", None)
with self.assertRaises(AssertionError):
self.network.perform_action(exploit)
# invalid service
exploit = Action((DMZ, 0), A_COST, "exploit", self.S + 1)
with self.assertRaises(AssertionError):
self.network.perform_action(exploit)
def test_step_exploit_user(self):
t_action = Action(self.ads_space[0], Environment.EXPLOIT_COST, "exploit", 0)
expected_obs = self.env.reset()
self.env.step(t_action)
self.update_obs(t_action, expected_obs, True)
t_action2 = Action(self.ads_space[2], Environment.EXPLOIT_COST, "exploit", 0)
o, r, d = self.env.step(t_action2)
self.update_obs(t_action2, expected_obs, True)
self.assertEqual(r, Environment.R_USER - t_action.cost)
self.assertFalse(d)
self.assertEqual(o, expected_obs)
def __init__(self, config, scan_cost=SCAN_COST, seed=1):
"""
Construct a new environment and network
Arguments:
dict config : network configuration
float scan_cost : cost of performing a scan action
int seed : random seed
"""
self.config = config
self.seed = seed
np.random.seed(seed)
self.num_services = config["num_services"]
self.network = Network(config)
self.address_space = self.network.get_address_space()
self.service_map = {}
self.service_exploits = config["service_exploits"]
for i, service in enumerate(self.service_exploits.keys()):
self.service_map[service] = i
self.action_space = Action.load_action_space(self.address_space, self.service_exploits,
scan_cost)
self.init_state = self._generate_initial_state()
self.compromised_subnets = None
self.renderer = None
self.reset()
def test_step_done(self):
t_action = Action(self.ads_space[0], Environment.EXPLOIT_COST, "exploit", 0)
expected_obs = self.env.reset()
self.env.step(t_action)
self.update_obs(t_action, expected_obs, True)
t_action1 = Action(self.ads_space[1], Environment.EXPLOIT_COST, "exploit", 0)
o, r, d = self.env.step(t_action1)
self.update_obs(t_action1, expected_obs, True)
t_action2 = Action(self.ads_space[2], Environment.EXPLOIT_COST, "exploit", 0)
o, r, d = self.env.step(t_action2)
self.update_obs(t_action2, expected_obs, True)
self.assertTrue(d)
self.assertEqual(o, expected_obs)
def test_step_not_reachable(self):
t_action = Action(self.ads_space[1], Environment.SCAN_COST, "scan")
expected_obs = self.env.reset()
o, r, d = self.env.step(t_action)
self.assertEqual(r, -t_action.cost)
self.assertFalse(d)
self.assertEqual(o, expected_obs)
def test_reset2(self):
t_action = Action(self.ads_space[0], Environment.SCAN_COST, "scan")
o, _, _ = self.env.step(t_action)
self.env.reset()
actual_obs = self.env.reset()
expected_obs = self.get_initial_expected_obs()
self.assertEqual(actual_obs, expected_obs)
def test_successful_exploit(self):
rewards = [0, self.r_sensitive, self.r_user]
for i in range(DMZ, USER + 1):
subnet = i
e = self.find_exploit(self.network, (subnet, 0), True)
exploit = Action((subnet, 0), A_COST, "exploit", e)
outcome, reward, services = self.network.perform_action(exploit)
self.assertTrue(outcome)
self.assertEqual(reward, rewards[i - 1])
def test_gen_action_space_mixed(self):
actual_A_space = Action.generate_action_space(self.ads_space, self.E,
self.test_exploit_cost, self.test_scan_cost,
exploit_probs="mixed")
for a in actual_A_space:
if a.is_scan():
self.assertTrue(a.prob == 1.0)
else:
self.assertTrue(a.prob == 0.5 or a.prob == 0.8)
actual_A_space = Action.generate_action_space(self.ads_space, 5,
self.test_exploit_cost, self.test_scan_cost,
exploit_probs="mixed")
for a in actual_A_space:
if a.is_scan():
self.assertTrue(a.prob == 1.0)
else:
self.assertTrue(a.prob == 0.2 or a.prob == 0.5 or a.prob == 0.8)
def test_scan(self):
rewards = [0, 0, 0]
for i in range(DMZ, USER + 1):
subnet = i
exp_services = self.network.machines[(subnet, 0)]._services
scan = Action((subnet, 0), A_COST, "scan", None)
outcome, reward, services = self.network.perform_action(scan)
self.assertTrue(outcome)
self.assertEqual(reward, rewards[i - 1])
self.assertTrue(services == exp_services)
def test_unsuccessful_exploit(self):
rewards = [0, 0, 0]
exp_services = np.asarray([])
for i in range(DMZ, USER + 1):
subnet = i
e = self.find_exploit(self.network, (subnet, 0), False)
if (e is None):
# machine vulnerable to all exploits
continue
exploit = Action((subnet, 0), A_COST, "exploit", e)
outcome, reward, services = self.network.perform_action(exploit)
self.assertFalse(outcome)
self.assertEqual(reward, rewards[i - 1])
self.assertTrue((services == exp_services).all())
def main():
"""
Test rendering of a single episode using Viewer class in render module
"""
generated_env = True
config_file = small_config
actions = [Action((0, 0), EXPLOIT_COST, "exploit", 0),
Action((1, 0), EXPLOIT_COST, "exploit", 0),
Action((2, 1), EXPLOIT_COST, "exploit", 0),
Action((2, 0), EXPLOIT_COST, "exploit", 0)]
if generated_env:
E = 1
M = 40
env = NetworkAttackSimulator.from_params(M, E)
else:
env = NetworkAttackSimulator.from_file(config_file)
test_render_asci(env, actions)
test_render_readable(env, actions)
test_render_network_graph(env)
test_render_episode(env, actions)
def test_scan(self):
# address is ignored as this is controlled and checked at Network level
exploit = Action((1, 1), A_COST, "scan", None)
# Test on machine with no sensitive docs (i.e. 0 value)
outcome, reward, services = self.test_m1.perform_action(exploit)
self.assertTrue(outcome)
self.assertEqual(reward, 0)
self.assertTrue((services == self.services).all())
# Test exploit on machine with sensitive docs
outcome, reward, services = self.test_m2.perform_action(exploit)
self.assertTrue(outcome)
self.assertEqual(reward, 0)
self.assertTrue((services == self.services).all())
def test_unsuccessful_exploit(self):
# address is ignored as this is controlled and checked at Network level
exploit = Action((1, 1), A_COST, "exploit", 1)
# Test on machine with no sensitive docs (i.e. 0 value)
outcome, reward, services = self.test_m1.perform_action(exploit)
self.assertFalse(outcome)
self.assertEqual(reward, 0)
self.assertTrue((services == np.asarray([])).all())
# Test exploit on machine with sensitive docs
outcome, reward, services = self.test_m2.perform_action(exploit)
self.assertFalse(outcome)
self.assertEqual(reward, 0)
self.assertTrue((services == np.asarray([])).all())