def main():
if len(sys.argv) != 3:
print("Usage: python demo_visualize_graph.py <scenario> <generate>")
return 1
scenario_name = sys.argv[1]
generate = bool(int(sys.argv[2]))
print(generate)
print("Displaying {} scenario".format(scenario_name))
if generate:
print("Generating network configuration")
scenario = get_scenario(scenario_name)
if scenario is None:
return 1
num_machines = scenario["machines"]
num_services = scenario["services"]
restrictiveness = scenario["restrictiveness"]
print("\tnumber of machines =", num_machines)
print("\tnumber of services =", num_services)
print("\tfirewall restrictiveness =", restrictiveness)
env = NetworkAttackSimulator.from_params(
num_machines, num_services, restrictiveness=restrictiveness)
else:
print("Loading network configuration")
env = NetworkAttackSimulator.from_file(scenario_name)
print("Max score:", env.get_best_possible_score())
env.render_network_graph(show=True)
def get_scenario_env(scenario_name, seed=1):
if not is_valid_scenario(scenario_name, verbose=True):
return None
scenario_params = scenarios[scenario_name]
if scenario_params["generate"]:
M = scenario_params["machines"]
S = scenario_params["services"]
rve = scenario_params["restrictiveness"]
env = Cyber.from_params(M, S,
r_sensitive=R_SENS, r_user=R_USR,
exploit_cost=COST_EXP, scan_cost=COST_SCAN,
restrictiveness=rve, exploit_probs=EXPLOIT_PROB, seed=seed)
else:
env = Cyber.from_file(scenario_params["file"], scan_cost=COST_SCAN, seed=seed)
return env, scenario_params
def run_scaling_experiment(agent_name, eval_file):
print("\nRunning scaling analysis for agent: \n\t {0}".format(
str(agent_name)))
for m in range(MACHINE_MIN, MACHINE_MAX + 1, MACHINE_INTERVAL):
for s in range(SERVICE_MIN, SERVICE_MAX + 1, SERVICE_INTERVAL):
print("\n>> Machines={0} Services={1}".format(m, s))
for t in range(START_SEED, START_SEED + SCALING_RUNS):
print("\tRun {0} of {1} seed={2}".format(
t + 1, SCALING_RUNS + START_SEED, t))
env = Cyber.from_params(m,
s,
r_sensitive=R_SENS,
r_user=R_USR,
exploit_cost=COST_EXP,
scan_cost=COST_SCAN,
restrictiveness=RVE,
exploit_probs=EXPLOIT_PROB,
seed=t)
print("Agent name")
agent = get_agent(agent_name, "default", env)
if agent_name != "random":
ep_tsteps, ep_rews, ep_times = agent.train(
env, MAX_EPISODES, MAX_STEPS, TIMEOUT, VERBOSE)
path_found = is_path_found(ep_tsteps, MAX_STEPS)
exp_reward = get_expected_rewards(ep_rews)
training_time = sum(ep_times)
print(
"\tTraining run {} - path_found={} - exp_reward={} - train_time={:.2f}"
.format(t, path_found, exp_reward, training_time))
run_evaluation(agent, env, (m, s), agent_name, MAX_STEPS, t,
eval_file)
def run_agent(scenario, agent_type, agent_params):
"""
Run agent on given environment for set number of runs and return averaged
rewards (averaged over runs)
"""
scenario_params = scenarios[scenario]
M = scenario_params["machines"]
S = scenario_params["services"]
rve = scenario_params["restrictiveness"]
num_episodes = scenario_params["episodes"]
max_steps = scenario_params["steps"]
timeout = scenario_params["timeout"]
run_rewards = []
for t in range(RUNS):
env = Cyber.from_params(M,
S,
r_sensitive=R_SENS,
r_user=R_USR,
exploit_cost=COST_EXP,
scan_cost=COST_SCAN,
restrictiveness=rve,
exploit_probs=EXPLOIT_PROB,
seed=t)
agent = load_agent(agent_type, env, agent_params)
ep_tsteps, ep_rews, ep_times = agent.train(env, num_episodes,
max_steps, timeout, VERBOSE)
run_rewards.append(ep_rews)
print("\t\tRun {} - total reward = {}".format(t, sum(ep_rews)))
return np.mean(run_rewards, axis=0)
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 pompdp_from_params(num_machines,
num_services,
exploit_probs=1.0,
uniform=False,
restrictiveness=5,
discount=0.95):
"""
Generate a pomdp from given set of environment parameters
"""
env = NetworkAttackSimulator.from_params(num_machines,
num_services,
exploit_probs=exploit_probs,
uniform=uniform,
restrictiveness=restrictiveness)
output_name = "{}_{}_{}_{}_{}.pomdp".format(num_machines, num_services,
exploit_probs, uniform,
restrictiveness)
generate_pomdp(env, output_name, discount)
def load_env(M, S, seed):
return Cyber.from_params(M, S, restrictiveness=RVE, exploit_probs=EXPLOIT_PROB, seed=seed)
def get_env(self, nM, nS):
env = NetworkAttackSimulator.from_params(nM, nS)
return env
def pomdp_from_file(config_path, output_name="pomdp.pomdp", discount=0.95):
"""
Generate a pomdp from a given config file
"""
env = NetworkAttackSimulator.from_file(config_path)
generate_pomdp(env, output_name, discount)
def main():
if len(sys.argv) != 4 and len(sys.argv) != 5:
print("Usage: python demo_solving.py agent scenario generate [seed]")
return 1
agent_name = sys.argv[1]
if not is_valid_agent(agent_name, verbose=True):
return 1
scenario_name = sys.argv[2]
generate = bool(int(sys.argv[3]))
if len(sys.argv) == 5:
seed = int(sys.argv[4])
else:
seed = 1
print("Displaying {} scenario".format(scenario_name))
if generate:
print("Generating network configuration")
scenario = get_scenario(scenario_name)
if scenario is None:
return 1
num_machines = scenario["machines"]
num_services = scenario["services"]
rve = scenario["restrictiveness"]
num_episodes = scenario["episodes"]
max_steps = scenario["steps"]
timeout = scenario["timeout"]
print("\tnumber of machines =", num_machines)
print("\tnumber of services =", num_services)
print("\tfirewall restrictiveness =", rve)
env = NetworkAttackSimulator.from_params(num_machines, num_services, restrictiveness=rve,
seed=seed)
agent_scenario = scenario_name
else:
print("Loading network configuration")
num_episodes = 200
max_steps = 500
timeout = 120
# env = CyberAttackSimulatorEnv.from_file(scenario_name)
env = NetworkAttackSimulator.from_params(51, 6,
r_sensitive=10, r_user=10,
exploit_cost=1, scan_cost=1,
restrictiveness=3, exploit_probs=0.7,
seed=2)
agent_scenario = "default"
agent = get_agent(agent_name, agent_scenario, env)
if agent is None:
return 1
print("Solving {} scenario using {} agent".format(scenario_name, agent_name))
train_args = {}
train_args["visualize_policy"] = num_episodes // 1
ep_tsteps, ep_rews, ep_times = agent.train(env, num_episodes, max_steps, timeout,
verbose=True, **train_args)
gen_episode = agent.generate_episode(env, max_steps)
env.render_episode(gen_episode)
plot_results(ep_tsteps, ep_rews, ep_times, env)