working on bugs

deploy.py

@@ -18,6 +18,42 @@ from datetime import datetime
 debug_mode = True
 deployment_id = None
 
+# Constants for providers
+PROVIDERS = ["aws", "linode", "flokinet"]
+DEFAULT_SSH_USER = {
+    "aws": "root",
+    "linode": "root",
+    "flokinet": "root"
+}
+
+# Directory names - maintain correct case for each provider
+PROVIDER_DIRS = {
+    "aws": "AWS",
+    "linode": "Linode",
+    "flokinet": "FlokiNET"
+}
+
+# Then fix the select_provider function
+def select_provider():
+    """Let the user select a cloud provider"""
+    print("\nAvailable cloud providers:")
+    for i, provider in enumerate(PROVIDERS, 1):
+        print(f"  {i}. {provider.capitalize()}")
+    
+    while True:
+        try:
+            provider_choice = input("\nSelect a provider (1-3 or 99 to cancel): ")
+            if provider_choice == "99":
+                return None
+            
+            provider_choice = int(provider_choice)
+            if 1 <= provider_choice <= len(PROVIDERS):
+                return PROVIDERS[provider_choice - 1]
+            else:
+                print(f"{COLORS['RED']}Please enter a number between 1 and {len(PROVIDERS)}{COLORS['RESET']}")
+        except ValueError:
+            print(f"{COLORS['RED']}Please enter a valid number{COLORS['RESET']}")
+
 # Color codes for terminal output
 COLORS = {
     "RESET": "\033[0m",
@@ -53,8 +89,12 @@ def print_banner():
 
 def main_menu():
     """Display the main menu and handle user selection"""
-    global debug_mode
+    global debug_mode, deployment_id
+    
     while True:
+        # Reset deployment ID for each new operation from the menu
+        deployment_id = generate_deployment_id()
+        
         clear_screen()
         print_banner()
         print(f"{COLORS['WHITE']}MAIN MENU{COLORS['RESET']}")
@@ -72,7 +112,7 @@ def main_menu():
         print(f"11) Custom Deployment")
         print(f"12) Tools")
         print(f"13) Debug Mode: {COLORS['GREEN'] if debug_mode else COLORS['RED']}{debug_mode}{COLORS['RESET']}")
-        print(f"")
+        print(f"\n")
         print(f"99) Exit")
         
         choice = input("\nSelect an option: ")
@@ -229,21 +269,24 @@ def deploy_tracker():
 
 def custom_deployment():
     """Run the full interactive deployment wizard"""
-    deployment_id = generate_deployment_id()
+    current_deployment_id = generate_deployment_id()
-    config = interactive_setup(deployment_id)
+    config = interactive_setup(current_deployment_id)
     if config:
         execute_deployment(config)
 
 def initialize_deployment():
-    """Initialize global deployment ID"""
+    """Initialize and return a fresh deployment ID"""
-    global deployment_id
+    new_deployment_id = generate_deployment_id()
-    deployment_id = generate_deployment_id()
+    logging.info(f"Initialized new deployment ID: {new_deployment_id}")
-    return deployment_id
+    return new_deployment_id
 
 def gather_common_parameters():
     """Collect common parameters needed for deployments"""
-    global debug_mode, deployment_id
+    global debug_mode
-    config = {'deployment_id': deployment_id}
+    
+    # Generate a fresh deployment ID for this specific deployment
+    current_deployment_id = generate_deployment_id()
+    config = {'deployment_id': current_deployment_id}
     config['debug'] = debug_mode
 
     # Get provider
@@ -317,13 +360,13 @@ def gather_common_parameters():
     config['letsencrypt_email'] = email
     
     # Set SSH key
-    config['ssh_key'] = generate_ssh_key(deployment_id)
+    config['ssh_key'] = generate_ssh_key(current_deployment_id)
     config['ssh_key_path'] = f"{config['ssh_key']}.pub"
     
     # Set consistent resource names based on deployment ID
-    config['redirector_name'] = f"r-{deployment_id}"
+    config['redirector_name'] = f"r-{current_deployment_id}"
-    config['c2_name'] = f"s-{deployment_id}"
+    config['c2_name'] = f"s-{current_deployment_id}"
-    config['tracker_name'] = f"t-{deployment_id}"
+    config['tracker_name'] = f"t-{current_deployment_id}"
     
     # Security options
     print("\nSecurity options:")
@@ -331,28 +374,11 @@ def gather_common_parameters():
     config['secure_memory'] = input("Enable secure memory settings? (y/n) [default: y]: ").lower() != 'n'
     config['zero_logs'] = input("Enable zero-logs configuration? (y/n) [default: y]: ").lower() != 'n'
     
+    # Add SSH option that was missing
+    config['ssh_after_deploy'] = input("\nSSH into instance after deployment? (y/n) [default: n]: ").lower() == 'y'
+    
     return config
 
-def select_provider():
-    """Let the user select a cloud provider"""
-    print("\nAvailable cloud providers:")
-    for i, provider in enumerate(PROVIDERS, 1):
-        print(f"  {i}. {provider.capitalize()}")
-    
-    while True:
-        try:
-            provider_choice = input("\nSelect a provider (1-3 or 99 to cancel): ")
-            if provider_choice == "99":
-                return None
-            
-            provider_choice = int(provider_choice)
-            if 1 <= provider_choice <= len(PROVIDERS):
-                return PROVIDERS[provider_choice - 1]
-            else:
-                print(f"{COLORS['RED']}Please enter a number between 1 and {len(PROVIDERS)}{COLORS['RESET']}")
-        except ValueError:
-            print(f"{COLORS['RED']}Please enter a valid number{COLORS['RESET']}")
-
 def get_aws_credentials(provider_vars):
     """Get AWS credentials from user or vars file"""
     default_aws_key = provider_vars.get('aws_access_key', '')
@@ -484,29 +510,16 @@ def execute_deployment(config):
     
     if success:
         print(f"\n{COLORS['GREEN']}Deployment completed successfully!{COLORS['RESET']}")
+        
+        # Explicitly handle SSH after deployment if requested
+        if config.get('ssh_after_deploy', False):
+            print(f"\n{COLORS['BLUE']}Connecting to instance via SSH...{COLORS['RESET']}")
+            ssh_to_instance(config)
     else:
         print(f"\n{COLORS['RED']}Deployment failed.{COLORS['RESET']}")
     
     input("\nPress Enter to return to menu...")
 
-# Disable Ansible host key checking
-os.environ["ANSIBLE_HOST_KEY_CHECKING"] = "False"
-
-# Constants for providers
-PROVIDERS = ["aws", "linode", "flokinet"]
-DEFAULT_SSH_USER = {
-    "aws": "root",
-    "linode": "root",
-    "flokinet": "root"
-}
-
-# Directory names - maintain correct case for each provider
-PROVIDER_DIRS = {
-    "aws": "AWS",
-    "linode": "Linode",
-    "flokinet": "FlokiNET"
-}
-
 def generate_random_string(length=8):
     """Generate a random string of letters and digits."""
     return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))
@@ -1253,6 +1266,7 @@ def deploy_infrastructure(config):
     provider = config['provider']
     logging.info(f"Deploying {provider} infrastructure...")
     
+    try:
         # Set provider-specific environment variables
         if provider == "aws":
             if config.get('aws_access_key'):
@@ -1346,6 +1360,15 @@ def deploy_infrastructure(config):
                 config['c2_ip'] = c2_config['c2_ip']
         
         return True
+    except Exception as e:
+        logging.error(f"Deployment failed with error: {str(e)}")
+        if config.get('debug'):
+            import traceback
+            logging.error(traceback.format_exc())
+        
+        # Clean up any partial resources that were created
+        cleanup_resources(config, interactive=True)
+        return False
 
 def deploy_flokinet_redirector(config):
     """Deploy FlokiNET redirector separately"""
@@ -1601,6 +1624,26 @@ def cleanup_resources(config, interactive=True):
             print("You can clean them up later by running with --teardown")
             return False
     
+    # Clean up SSH keys
+    if 'deployment_id' in config:
+        ssh_key_path = f"~/.ssh/c2deploy_{config['deployment_id']}.pem"
+        expanded_path = os.path.expanduser(ssh_key_path)
+        if os.path.exists(expanded_path):
+            try:
+                os.remove(expanded_path)
+                logging.info(f"Removed SSH key: {ssh_key_path}")
+            except Exception as e:
+                logging.error(f"Failed to remove SSH key {ssh_key_path}: {e}")
+                
+        # Also check for public key
+        pub_key_path = f"{expanded_path}.pub"
+        if os.path.exists(pub_key_path):
+            try:
+                os.remove(pub_key_path)
+                logging.info(f"Removed SSH public key: {pub_key_path}.pub")
+            except Exception as e:
+                logging.error(f"Failed to remove SSH public key {pub_key_path}.pub: {e}")
+    
     # Use Ansible for cleanup with confirmation set to false
     extra_vars = {
         "confirm_cleanup": False,  # Skip confirmation prompt
@@ -1655,6 +1698,7 @@ def cleanup_resources(config, interactive=True):
             
     return True
 
+
 def check_dependencies():
     """Check if required dependencies are installed"""
     dependencies = {