Proxy set-up for CUBE

Python/ssh_fabric_batch.py

@@ -5,6 +5,10 @@ from dotenv import load_dotenv
 import io # Make sure io is imported at the top of your script
 import os
 from cube_activate_ssh import activate_ssh
+from ruamel.yaml import YAML
+from ruamel.yaml.scalarstring import DoubleQuotedScalarString
+import shlex
+import base64
 
 def execute_command(c, command):
     """Executes a simple command on the remote device."""
@@ -45,7 +49,7 @@ def read_remote_config_sudo(c, remote_path, sudo_pass):
         print(f"Error reading remote file with sudo: {e}")
         return None
 
-import shlex # Make sure to import shlex at the top of your script
+ # Make sure to import shlex at the top of your script
 
 def write_remote_config_sudo(c, remote_path, content, sudo_pass, user_owner, group_owner, permissions):
     """
@@ -261,11 +265,187 @@ def parse_connection_string(connection_string):
             
     return parsed_data
 
+def find_yaml_value(yaml_content, key_path):
+    """
+    Finds a value in a YAML string using a dot-separated path.
+
+    Args:
+        yaml_content (str): The string content of the YAML file.
+        key_path (str): A dot-separated path to the key (e.g., "cubeProcess.cyber_check").
+
+    Returns:
+        The value if found, otherwise None.
+    """
+    try:
+        yaml = YAML()
+        data = yaml.load(yaml_content)
+        
+        # Traverse the path
+        keys = key_path.split('.')
+        current_level = data
+        for key in keys:
+            current_level = current_level[key]
+        
+        return current_level
+    except (KeyError, TypeError):
+        # KeyError if a key is not found, TypeError if trying to index a non-dict
+        # print(f"Warning: Key path '{key_path}' not found in YAML content.")
+        return None
+    
+def set_yaml_value(yaml_content, key_path, new_value):
+    """
+    Sets a value in a YAML string using a dot-separated path.
+    Preserves comments, formatting, and quotes thanks to ruamel.yaml.
+    This version correctly traverses nested keys.
+
+    Args:
+        yaml_content (str): The string content of the YAML file.
+        key_path (str): A dot-separated path to the key (e.g., "cubeProcess.cyber_check").
+        new_value: The new value to set.
+
+    Returns:
+        str: The modified YAML content as a string, or the original content on error.
+    """
+    try:
+        # --- FIX 1: Configure the YAML object to preserve quotes ---
+        yaml = YAML()
+        yaml.preserve_quotes = True
+        yaml.indent(mapping=2, sequence=4, offset=2) # Optional: ensures consistent indentation
+
+        data = yaml.load(yaml_content)
+        
+        # --- FIX 2: Correct traversal logic ---
+        keys = key_path.split('.')
+        current_level = data
+        
+        # Traverse down to the final key's parent dictionary
+        for key in keys[:-1]:
+            current_level = current_level[key]
+            
+        final_key = keys[-1]
+        
+        # Check if the key exists before setting it
+        if final_key not in current_level:
+            print(f"❌ Error: Final key '{final_key}' not found in the structure. Aborting.")
+            return yaml_content # Return original content
+
+        # Set the new value
+        current_level[final_key] = new_value
+        
+        # Dump the modified data back to a string
+        string_stream = io.StringIO()
+        yaml.dump(data, string_stream)
+        return string_stream.getvalue()
+        
+    except (KeyError, TypeError) as e:
+        print(f"❌ Error: Key path '{key_path}' is invalid or part of the path does not exist. Error: {e}")
+        return yaml_content # Return original content on failure
+def ensure_iptables_port_rule(config_content, target_port, template_port):
+    """
+    Ensures that iptables rules for a target port exist in the configuration.
+
+    If rules for the target port are not found, it finds rules for a
+    template port and replaces the port number.
+
+    Args:
+        config_content (str): The multi-line string of the iptables rules file.
+        target_port (int or str): The port number that should exist (e.g., 8080).
+        template_port (int or str): The port number to use as a template (e.g., 443).
+
+    Returns:
+        str: The modified (or original) configuration content.
+    """
+    target_port_str = str(target_port)
+    template_port_str = str(template_port)
+    lines = config_content.splitlines()
+    
+    target_port_found = False
+    
+    # --- PASS 1: Check if the target port rule already exists ---
+    for line in lines:
+        # Check for the target port in an active rule line
+        # The spaces around the port string prevent accidentally matching '8080' in '18080'
+        if line.strip().startswith('-A') and (f"--dport {target_port_str}" in line or f"--sport {target_port_str}" in line):
+            print(f"✅ Info: Rule for target port {target_port_str} already exists. No changes needed.")
+            target_port_found = True
+            break
+            
+    # If the rule was found, return the original content without any changes.
+    if target_port_found:
+        return config_content
+        
+    # --- PASS 2: If we get here, the rule was not found. We must replace the template. ---
+    print(f"Info: Rule for target port {target_port_str} not found. Searching for template port {template_port_str} to replace.")
+    
+    new_lines = []
+    changes_made = False
+    for line in lines:
+        # Check for the template port in an active rule line
+        if line.strip().startswith('-A') and (f"--dport {template_port_str}" in line or f"--sport {template_port_str}" in line):
+            # This is a line we need to modify
+            modified_line = line.replace(template_port_str, target_port_str)
+            new_lines.append(modified_line)
+            print(f"   - Replacing: '{line}'")
+            print(f"   + With:      '{modified_line}'")
+            changes_made = True
+        else:
+            # This line doesn't need changing, add it as is.
+            new_lines.append(line)
+            
+    if not changes_made:
+        print(f"❌ Warning: Target port {target_port_str} was not found, AND template port {template_port_str} was also not found. No changes made.")
+        return config_content # Return original if template wasn't found either
+
+    return "\n".join(new_lines)
+
+def write_remote_config_base64_sudo(c, remote_path, content, sudo_pass, user_owner, group_owner, permissions):
+    """
+    Writes content directly to a remote file by passing it as a Base64 string.
+
+    This is the most robust method for no-SFTP environments, as it completely
+    avoids all shell quoting and parsing issues for complex, multi-line content.
+
+    Args:
+        c (fabric.Connection): The active connection object.
+        remote_path (str): The absolute path to the file on the remote host.
+        content (str): The string content to be written to the file.
+        sudo_pass (str): The sudo password for the write operation.
+    """
+    print(f"\n--- [{c.host}] Writing content via Base64 to: {remote_path} ---")
+    try:
+        # Step 1: Encode the string content into Base64.
+        # base64.b64encode requires bytes, so we encode the string to utf-8.
+        # The result is bytes, so we decode it back to a simple ascii string to use in our command.
+        base64_content = base64.b64encode(content.encode('utf-8')).decode('ascii')
+
+        # Step 2: Construct the command.
+        # 'echo ... | base64 --decode > file': This pipeline decodes the content
+        # and redirects the output to the destination file.
+        # We wrap the entire pipeline in 'sudo sh -c "..."' so that the
+        # redirection ('>') is performed by a shell running as root.
+        command = f"sh -c \"echo '{base64_content}' | base64 --decode > {remote_path}\""
+
+        print("Step 1: Writing Base64 content via root shell...")
+        c.sudo(command, password=sudo_pass, hide=True)
+
+        # Step 3: Set ownership and permissions.
+        print("Step 2: Setting ownership and permissions...")
+        c.sudo(f"chown {user_owner}:{group_owner} {remote_path}", password=sudo_pass)
+        c.sudo(f"chmod {permissions} {remote_path}", password=sudo_pass)
+
+        print(f"✅ Successfully wrote content to {remote_path}")
+
+    except Exception as e:
+        print(f"❌ Error writing Base64 content to remote file: {e}")
+        # Re-raise the exception for the main loop.
+        raise
+
+
 def main():
     """Main function to parse arguments and orchestrate tasks."""
     ip_address_prefix = "10.81.35." # Carling subnet
-    ip_address_range = [] # list(range(65, 75)) # From 65 to 74
+    ip_address_range = list(range(68, 75)) # Fronm 65 to 74
-    ip_address_range.append(66) #(85) # Add 85 after 74.
+    ip_address_range.append(85) # Add 85 after 74.
     hosts = [f"{ip_address_prefix}{suffix}" for suffix in ip_address_range]
 
     ssh_port = 11022
@@ -305,7 +485,7 @@ def main():
                 print(f"Checking Cloud configuration:", end=" ", flush=True)
                 result = read_remote_config_sudo(c, "/etc/cube/config-azure.properties", ssh_password)
                 print(f"✅", end="\n", flush=True)
-            except:
+            except Exception as e:
                 print(f"❌", end="\n", flush=True)
                 print(f"[Cloud configuration check] Exception: {e}")
                 continue
@@ -318,7 +498,7 @@ def main():
             result = result_proxy_host_port
             cloud_configuration_check(hostname, result, "iot-ingest-ess-prod.azure-devices.net", "10.81.35.126", "8080")
 
-            response = input(f"Apply the change on {hostname.strip()}? (y)es to apply, anything else to cancel - ").lower()
+            response = input(f"Apply the change on {hostname.strip()}? (y)es or (n)o, anything else to cancel - ").lower()
             if response in ['y']:
                 print(f"Applying changes:", end=" ", flush=True)
                 try:
@@ -332,15 +512,141 @@ def main():
                 try:
                     result = read_remote_config_sudo(c, "/etc/cube/config-azure.properties", ssh_password)
                     print(f"✅", end="\n", flush=True)
-                except:
+                except Exception as e:
                     print(f"❌", end="\n", flush=True)
                     print(f"[Proxy verification] Exception: {e}")
                     continue
 
                 cloud_configuration_check(hostname, result, "iot-ingest-ess-prod.azure-devices.net", "10.81.35.126", "8080")
+            elif response in ['n']:
+                print(f"Not applying configuration...")
             else:
                 print(f"Not applying configuration...")
                 continue
 
+            print(f"Disabling Cyber Check:", end=" ", flush=True)
+            try:
+                execute_sudo_command(c, "systemctl stop cube-monit.service", ssh_password)
+                execute_sudo_command(c, "mount -o remount,rw /", ssh_password)
+                print(f"✅", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[Disabling Cyber Check] Exception: {e}")
+                continue
+
+            print(f"Reading Cyber Check configuration:", end=" ", flush=True)
+            try:
+                result = read_remote_config_sudo(c, "/etc/cube-default/configfile_monit.yaml", ssh_password)
+                print(f"✅", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[Cyber Check configuration] Exception: {e}")
+                continue
+
+            print(f"Checking cyber_check:", end=" ", flush=True)
+            try:
+                status = find_yaml_value(result, "cubeProcess.cyber_check")
+                if status == False:
+                    print(f"✅", end="\n", flush=True)
+                else:
+                    print(f"❌", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[cyber_check value] Exception: {e}")
+                continue
+
+            print(f"Modifying cyber_check:", end=" ", flush=True)
+            modified_result = ""
+            try:
+                modified_result = set_yaml_value(result, "cubeProcess.cyber_check", False)
+                print(f"✅", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[cyber_check modification] Exception: {e}")
+                continue
+
+            print(f"Checking modified cyber_check:", end=" ", flush=True)
+            try:
+                status = find_yaml_value(modified_result, "cubeProcess.cyber_check")
+                if status == False:
+                    print(f"✅", end="\n", flush=True)
+                else:
+                    print(f"❌", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[Modified cyber_check value] Exception: {e}")
+                continue
+
+            response = input(f"Apply the change on {hostname.strip()}? (y)es or (n)o, anything else to cancel - ").lower()
+            if response in ['y']:
+                print(f"Applying changes:", end=" ", flush=True)
+                try:
+                    write_remote_config_base64_sudo(c,  "/etc/cube-default/configfile_monit.yaml", modified_result, ssh_password, "root", "root", "644")
+                    print(f"✅", end="\n", flush=True)
+                except Exception as e:
+                    print(f"❌", end="\n", flush=True)
+                    print(f"[cyber_check configuration] Exception: {e}")
+                    continue
+                print(f"Checking cyber_check configuration:", end=" ", flush=True)
+                try:
+                    result = read_remote_config_sudo(c, "/etc/cube-default/configfile_monit.yaml", ssh_password)
+                    print(f"✅", end="\n", flush=True)
+                except Exception as e:
+                    print(f"❌", end="\n", flush=True)
+                    print(f"[cyber_check configuration] Exception: {e}")
+                    continue
+                try:
+                    status = find_yaml_value(result, "cubeProcess.cyber_check")
+                    if status == False:
+                        print(f"✅", end="\n", flush=True)
+                    else:
+                        print(f"❌", end="\n", flush=True)
+                except Exception as e:
+                    print(f"❌", end="\n", flush=True)
+                    print(f"[Modified cyber_check configuration verification] Exception: {e}")
+                    continue
+            elif response in ['n']:
+                print(f"Not applying configuration...")
+            else:
+                print(f"Not applying configuration...")
+                continue
+
+
+            print(f"Firewall check:", end="\n", flush=True)
+            modified_result = ""
+            try:
+                result = read_remote_config_sudo(c, "/etc/iptables/iptables-cube.rules", ssh_password)
+            except Exception as e:
+                print(f"[Firewall reading] Exception: {e}")
+                continue
+            try:
+                modified_result = ensure_iptables_port_rule(result, 8080, 443)
+            except Exception as e:
+                print(f"[Firewall changes] Exception: {e}")
+                continue
+
+            response = input(f"Apply the change on {hostname.strip()}? (y)es or (n)o, anything else to cancel - ").lower()
+            if response in ['y']:
+                try:
+                    write_remote_config_base64_sudo(c, "/etc/iptables/iptables-cube.rules", modified_result, ssh_password, "root", "root", 600)
+                except Exception as e:
+                    print(f"[Firewall configuration] Exception: {e}")
+                    continue
+            elif response in ['n']:
+                print(f"Not applying configuration...")
+            else:
+                print(f"Not applying configuration...")
+                continue
+
+            print(f"Restarting Cyber Check:", end=" ", flush=True)
+            try:
+                execute_sudo_command(c, "mount -o remount,ro /", ssh_password)
+                execute_sudo_command(c, "systemctl start cube-monit.service", ssh_password)
+                print(f"✅", end="\n", flush=True)
+            except Exception as e:
+                print(f"❌", end="\n", flush=True)
+                print(f"[Restarting Cyber Check] Exception: {e}")
+                continue
+
 if __name__ == "__main__":
     main()