deploy.py

#!/usr/bin/env python3
"""
Deployment script for RocketGraph on AWS, Azure, or GCP.

Reads deployment parameters from a YAML or JSON profile file and uses them to
generate terraform.tfvars, then runs terraform init/plan/apply.

Usage:
    python deploy.py --profile config/default.yml
    python deploy.py --profile config/customer-a.yml --plan
    python deploy.py --profile config/customer-a.yml --destroy
    python deploy.py --list-profiles
"""
import os
import sys
import json
import argparse
import subprocess
from pathlib import Path

try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False


def _deep_merge(base: dict, override: dict) -> dict:
    """Recursively merge override into base; override values win on conflicts."""
    result = base.copy()
    for key, value in override.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = _deep_merge(result[key], value)
        else:
            result[key] = value
    return result


def _load_yaml_file(path: Path) -> dict:
    """Load a single YAML or JSON file and return its contents as a dict."""
    with open(path) as f:
        if path.suffix in ('.yml', '.yaml'):
            if not HAS_YAML:
                print("PyYAML is required: pip install pyyaml")
                sys.exit(1)
            return yaml.safe_load(f) or {}
        else:
            return json.load(f)


def _resolve_extends(ref: str, current_path: Path) -> Path:
    """Resolve an extends reference to an absolute path.

    A bare name like 'haglin' is looked up as haglin.yml in the same directory
    as the current profile.  A path with an extension is used as-is (relative
    to the current profile's directory if not absolute).
    """
    ref_path = Path(ref)
    if not ref_path.suffix:
        ref_path = current_path.parent / (ref + current_path.suffix)
    elif not ref_path.is_absolute():
        ref_path = current_path.parent / ref_path
    return ref_path


def _load_profile_chain(path: Path, visited: set) -> dict:
    """Recursively load a profile, its inheritance chain, and secrets sidecars.

    Resolution order (last write wins):
      grandparent.yml → grandparent.secrets.yml
        → parent.yml → parent.secrets.yml
          → child.yml → child.secrets.yml

    The 'extends' key is consumed and removed from the returned profile.
    Circular inheritance raises ValueError.
    """
    canonical = str(path.resolve())
    if canonical in visited:
        raise ValueError(f"Circular profile inheritance detected: {path.name}")
    visited.add(canonical)

    profile = _load_yaml_file(path)
    extends = profile.pop('extends', None)

    if extends:
        parent_path = _resolve_extends(extends, path)
        if not parent_path.exists():
            print(f"Extended profile not found: {parent_path}  (referenced from {path.name})")
            sys.exit(1)
        # Recursively resolve the parent chain (including the parent's own secrets)
        base = _load_profile_chain(parent_path, visited)
        profile = _deep_merge(base, profile)

    # Merge this level's secrets sidecar: foo.yml -> foo.secrets.yml
    secrets_path = path.with_name(path.stem + '.secrets' + path.suffix)
    if secrets_path.exists():
        secrets = _load_yaml_file(secrets_path)
        profile = _deep_merge(profile, secrets)
        print(f"  + {secrets_path.name}")

    return profile


def load_profile(profile_path: str) -> dict:
    """Load a deployment profile with inheritance and secrets sidecars.

    Profiles may declare 'extends: <name>' to inherit from another profile.
    Each level's secrets sidecar (foo.secrets.yml) is merged immediately after
    its profile, so child profiles can override parent secrets if needed.

    See config/default.secrets.yml.example for the secrets sidecar format.
    """
    path = Path(profile_path)
    if not path.exists():
        # Try resolving a bare name to config/<name>.yml
        candidate = Path(__file__).parent / 'config' / (profile_path + '.yml')
        if candidate.exists():
            path = candidate
        else:
            print(f"Profile not found: {profile_path}")
            sys.exit(1)
    print(f"Loading profile: {path}")
    return _load_profile_chain(path, set())


def list_profiles():
    """Print all invocable profile files found in config/.

    Secrets sidecars (*.secrets.yml) and example templates (*.example) are
    excluded — they are supporting files, not profiles you pass to --profile.
    """
    config_dir = Path(__file__).parent / 'config'
    candidates = (sorted(config_dir.glob('*.yml')) +
                  sorted(config_dir.glob('*.yaml')) +
                  sorted(config_dir.glob('*.json')))
    profiles = [p for p in candidates
                if '.secrets.' not in p.name and not p.name.endswith('.example')]
    if not profiles:
        print("No profiles found in config/")
        return
    print("Available profiles:")
    for p in profiles:
        # Show the extends chain so the user can see the inheritance structure
        try:
            data = _load_yaml_file(p)
            extends = data.get('extends')
            suffix = f"  (extends: {extends})" if extends else ""
        except Exception:
            suffix = ""
        print(f"  {p.relative_to(Path(__file__).parent)}{suffix}")


# SSM parameter definitions: profile key → (ssm path suffix, parameter type)
# SecureString parameters are masked in dry-run output.
_AWS_THREATWORX_PARAMS = [
    ('threatworx_api_token',        'api_token',        'SecureString'),
    ('threatworx_base_url',         'base_url',         'String'),
    ('threatworx_assets_base_url',  'assets_base_url',  'String'),
    ('threatworx_handle',           'handle',           'String'),
]

_AWS_XGT_PARAMS = [
    ('xgt_username',       'xgt_username',       'String'),
    ('xgt_password',       'xgt_password',       'SecureString'),
    ('xgt_security_label', 'xgt_security_label', 'String'),
]


def _mask(value: str) -> str:
    """Return the value with all but the first 4 characters masked."""
    if len(value) <= 4:
        return '****'
    return value[:4] + '*' * (len(value) - 4)


class CloudDeployer:
    def __init__(self, profile: dict):
        self.profile = profile
        self.provider = profile.get('provider', 'aws').lower()
        self.terraform_dir = Path(__file__).parent / 'terraform' / self.provider
        self.tfvars_file = self.terraform_dir / 'terraform.tfvars'

    def validate_provider(self):
        if self.provider not in ['aws', 'azure', 'gcp']:
            raise ValueError(f"Unsupported provider: {self.provider}. Must be aws, azure, or gcp.")
        if not self.terraform_dir.exists():
            raise FileNotFoundError(f"Terraform directory not found: {self.terraform_dir}")

    def check_prerequisites(self):
        """Check that required CLI tools are installed."""
        required = ['terraform']
        provider_tools = {'aws': ['aws'], 'azure': ['az'], 'gcp': ['gcloud']}
        required.extend(provider_tools.get(self.provider, []))

        missing = [t for t in required
                   if subprocess.run(['which', t], capture_output=True).returncode != 0]
        if missing:
            print(f"Missing required tools: {', '.join(missing)}")
            sys.exit(1)
        print("All prerequisites met")

    def validate_credentials(self):
        """Validate cloud provider credentials are active."""
        print(f"Validating {self.provider.upper()} credentials...")
        checks = {
            'aws':   (['aws', 'sts', 'get-caller-identity'],
                      "AWS credentials not configured. Run 'aws configure' first."),
            'azure': (['az', 'account', 'show'],
                      "Azure credentials not configured. Run 'az login' first."),
            'gcp':   (['gcloud', 'auth', 'application-default', 'print-access-token'],
                      "GCP credentials not configured. Run 'gcloud auth application-default login' first."),
        }
        cmd, msg = checks[self.provider]
        if subprocess.run(cmd, capture_output=True).returncode != 0:
            print(msg)
            sys.exit(1)
        print("Credentials validated")

    def push_secrets(self, dry_run: bool = False):
        """Push credentials from the profile into the cloud provider's secrets store.

        For AWS, writes each credential to SSM Parameter Store under the prefix
        defined by aws_secrets_prefix (default: /rocketworx/threatworx).
        Only parameters that are present and non-empty in the profile are written.

        With dry_run=True, prints the exact parameter names and values that would
        be stored (SecureString values are masked) without calling the cloud API.
        """
        if self.provider == 'aws':
            self._push_secrets_aws(dry_run)
        elif self.provider == 'azure':
            self._push_secrets_azure(dry_run)
        elif self.provider == 'gcp':
            self._push_secrets_gcp(dry_run)

    def _push_secrets_aws(self, dry_run: bool):
        p = self.profile
        tw_prefix  = p.get('threatworx_ssm_prefix', '/rocketworx/threatworx').rstrip('/')
        xgt_prefix = p.get('xgt_ssm_prefix',        '/rocketworx/xgt').rstrip('/')
        region     = p.get('region', 'us-west-2')

        # Collect ThreatWorx parameters present in the profile
        tw_params = []
        for profile_key, ssm_suffix, param_type in _AWS_THREATWORX_PARAMS:
            value = p.get(profile_key)
            if value:
                tw_params.append((f"{tw_prefix}/{ssm_suffix}", param_type, str(value)))

        # Collect XGT parameters present in the profile
        xgt_params = []
        for profile_key, ssm_suffix, param_type in _AWS_XGT_PARAMS:
            value = p.get(profile_key)
            if value:
                xgt_params.append((f"{xgt_prefix}/{ssm_suffix}", param_type, str(value)))

        # xgt_security_label defaults to xgt_username when not set explicitly
        if p.get('xgt_username') and not p.get('xgt_security_label'):
            xgt_params.append((f"{xgt_prefix}/xgt_security_label", 'String', str(p['xgt_username'])))

        if not tw_params and not xgt_params:
            print("No credentials found in profile — skipping SSM push.")
            print("Add credentials to your .secrets.yml sidecar.")
            return

        def _print_section(heading, prefix, params):
            if not params:
                return
            print(f"  {heading}  ({prefix})")
            col = max(len(name) for name, _, _ in params) + 2
            print(f"  {'Parameter Name':<{col}}  {'Type':<12}  Value")
            print(f"  {'-' * col}  {'-' * 12}  {'-' * 30}")
            for name, param_type, value in params:
                display = _mask(value) if param_type == 'SecureString' else value
                print(f"  {name:<{col}}  {param_type:<12}  {display}")
            print()

        def _put_params(params, label):
            print(f"Pushing {len(params)} {label} secret(s) to AWS SSM ...")
            for name, param_type, value in params:
                result = subprocess.run(
                    ['aws', 'ssm', 'put-parameter',
                     '--name', name,
                     '--value', value,
                     '--type', param_type,
                     '--overwrite',
                     '--region', region],
                    capture_output=True, text=True
                )
                if result.returncode == 0:
                    print(f"  ✅ {name}")
                else:
                    print(f"  ❌ {name}: {result.stderr.strip()}")
                    sys.exit(1)

        if dry_run:
            print("DRY RUN — no AWS calls will be made")
            print(f"  Region : {region}")
            print()
            _print_section("ThreatWorx", tw_prefix, tw_params)
            _print_section("XGT / RocketGraph", xgt_prefix, xgt_params)
            print("Run without --dry-run to store these values in AWS SSM.")
            return

        if tw_params:
            _put_params(tw_params, "ThreatWorx")
        if xgt_params:
            _put_params(xgt_params, "XGT / RocketGraph")
        print("Secrets pushed successfully.")

    def _push_secrets_azure(self, dry_run: bool):
        p = self.profile
        vault = p.get('azure_key_vault', '')
        if not vault:
            print("No azure_key_vault set in profile — skipping Key Vault push.")
            return

        mappings = [
            ('threatworx_api_token',       'threatworx-api-token'),
            ('threatworx_base_url',        'threatworx-base-url'),
            ('threatworx_assets_base_url', 'threatworx-assets-base-url'),
            ('threatworx_handle',          'threatworx-handle'),
            ('xgt_username',               'xgt-username'),
            ('xgt_password',               'xgt-password'),
            ('xgt_security_label',         'xgt-security-label'),
        ]
        params = [(secret_name, str(p[key])) for key, secret_name in mappings if p.get(key)]

        if dry_run:
            print(f"DRY RUN — no Azure calls will be made")
            print(f"  Vault : {vault}")
            print()
            for secret_name, value in params:
                print(f"  {secret_name}: {value}")
            return

        print(f"Pushing {len(params)} secret(s) to Azure Key Vault ({vault}) ...")
        for secret_name, value in params:
            result = subprocess.run(
                ['az', 'keyvault', 'secret', 'set',
                 '--vault-name', vault, '--name', secret_name, '--value', value],
                capture_output=True, text=True
            )
            if result.returncode == 0:
                print(f"  ✅ {secret_name}")
            else:
                print(f"  ❌ {secret_name}: {result.stderr.strip()}")
                sys.exit(1)

    def _push_secrets_gcp(self, dry_run: bool):
        p = self.profile
        project = p.get('gcp_secret_project', '')
        if not project:
            print("No gcp_secret_project set in profile — skipping Secret Manager push.")
            return

        mappings = [
            ('threatworx_api_token',       'threatworx-api-token'),
            ('threatworx_base_url',        'threatworx-base-url'),
            ('threatworx_assets_base_url', 'threatworx-assets-base-url'),
            ('threatworx_handle',          'threatworx-handle'),
            ('xgt_username',               'xgt-username'),
            ('xgt_password',               'xgt-password'),
            ('xgt_security_label',         'xgt-security-label'),
        ]
        params = [(secret_name, str(p[key])) for key, secret_name in mappings if p.get(key)]

        if dry_run:
            print(f"DRY RUN — no GCP calls will be made")
            print(f"  Project : {project}")
            print()
            for secret_name, value in params:
                print(f"  {secret_name}: {value}")
            return

        print(f"Pushing {len(params)} secret(s) to GCP Secret Manager ({project}) ...")
        for secret_name, value in params:
            result = subprocess.run(
                ['gcloud', 'secrets', 'versions', 'add', secret_name,
                 '--data-file=-', '--project', project],
                input=value, capture_output=True, text=True
            )
            if result.returncode == 0:
                print(f"  ✅ {secret_name}")
            else:
                print(f"  ❌ {secret_name}: {result.stderr.strip()}")
                sys.exit(1)

    def generate_tfvars(self):
        """Write terraform.tfvars from the profile. Overwrites any existing file."""
        p = self.profile
        print(f"Writing {self.tfvars_file} ...")

        if self.provider == 'aws':
            # key_name is required for AWS
            if not p.get('key_name'):
                print("Profile must include 'key_name' (AWS EC2 key pair name) for AWS deployments.")
                sys.exit(1)
            tfvars = {
                'aws_region':        p.get('region', 'us-west-2'),
                'project_name':      p.get('project_name', 'rocketgraph'),
                'environment':       p.get('environment', 'production'),
                'instance_type':     p.get('instance_type', 't4g.medium'),
                'key_name':          p['key_name'],
                'ssh_allowed_ips':   p.get('ssh_allowed_ips', ['0.0.0.0/0']),
                'root_volume_size':  p.get('disk_size', 30),
            }
            for key in ('instance_name', 'compose_url', 'env_url',
                        'threatworx_ssm_prefix', 'xgt_ssm_prefix',
                        'domain_name', 'ssl_email'):
                if p.get(key):
                    tfvars[key] = p[key]

        elif self.provider == 'azure':
            tfvars = {
                'project_name':        p.get('project_name', 'rocketgraph'),
                'location':            p.get('region', 'East US'),
                'vm_size':             p.get('instance_type', 'Standard_B2ms'),
                'ssh_public_key_path': p.get('ssh_public_key_path', '~/.ssh/id_rsa.pub'),
                'os_disk_size':        p.get('disk_size', 30),
            }
            for key in ('compose_url', 'env_url', 'azure_key_vault'):
                if p.get(key):
                    tfvars[key] = p[key]

        elif self.provider == 'gcp':
            if not p.get('gcp_project_id'):
                print("Profile must include 'gcp_project_id' for GCP deployments.")
                sys.exit(1)
            tfvars = {
                'project_id':          p['gcp_project_id'],
                'project_name':        p.get('project_name', 'rocketgraph'),
                'region':              p.get('region', 'us-central1'),
                'zone':                p.get('zone', 'us-central1-a'),
                'machine_type':        p.get('instance_type', 'e2-medium'),
                'ssh_public_key_path': p.get('ssh_public_key_path', '~/.ssh/id_rsa.pub'),
                'ssh_username':        p.get('ssh_username', 'ubuntu'),
                'ssh_allowed_ips':     p.get('ssh_allowed_ips', ['0.0.0.0/0']),
                'boot_disk_size':      p.get('disk_size', 30),
            }
            for key in ('compose_url', 'env_url', 'gcp_secret_project'):
                if p.get(key):
                    tfvars[key] = p[key]

        with open(self.tfvars_file, 'w') as f:
            for key, value in tfvars.items():
                if value is None or value == '':
                    continue
                if isinstance(value, str):
                    f.write(f'{key} = "{value}"\n')
                elif isinstance(value, list):
                    f.write(f'{key} = {json.dumps(value)}\n')
                elif isinstance(value, bool):
                    f.write(f'{key} = {str(value).lower()}\n')
                else:
                    f.write(f'{key} = {value}\n')

        print(f"Variables written to {self.tfvars_file}")

    def run_terraform(self, action='apply'):
        """Run terraform init and then the requested action."""
        os.chdir(self.terraform_dir)

        print("\nInitializing Terraform...")
        if subprocess.run(['terraform', 'init'], capture_output=False).returncode != 0:
            print("Terraform init failed")
            sys.exit(1)

        auto_approve = self.profile.get('auto_approve', False)

        if action == 'plan':
            print("\nRunning Terraform plan...")
            subprocess.run(['terraform', 'plan'])

        elif action == 'apply':
            print("\nApplying Terraform configuration...")
            cmd = ['terraform', 'apply'] + (['-auto-approve'] if auto_approve else [])
            if subprocess.run(cmd).returncode == 0:
                print("\nDeployment successful!")
                self.show_outputs()
            else:
                print("Deployment failed")
                sys.exit(1)

        elif action == 'destroy':
            print("\nDestroying infrastructure...")
            cmd = ['terraform', 'destroy'] + (['-auto-approve'] if auto_approve else [])
            subprocess.run(cmd)

    def show_outputs(self):
        """Display Terraform outputs after a successful apply."""
        print("\n" + "=" * 50)
        print("DEPLOYMENT INFORMATION")
        print("=" * 50)

        result = subprocess.run(['terraform', 'output', '-json'],
                                capture_output=True, text=True)
        if result.returncode != 0:
            return

        try:
            outputs = json.loads(result.stdout)
        except json.JSONDecodeError:
            return

        ip = outputs.get('public_ip', {}).get('value')
        if ip:
            print(f"\nPublic IP: {ip}")
            print(f"RocketGraph URL: http://{ip}:8080")

        if 'ssh_command' in outputs:
            print(f"\nSSH: {outputs['ssh_command']['value']}")

        if ip and self.provider == 'aws':
            key = self.profile.get('key_name', 'your-key')
            print("\nNext steps:")
            print("  1. Wait 10-20 minutes for installation to complete")
            print(f"  2. SSH in:     ssh -i ~/.ssh/{key}.pem ubuntu@{ip}")
            print("  3. Check logs: sudo tail -f /var/log/user-data.log")

        print(f"\nTo destroy: python deploy.py --profile <profile> --destroy")


def main():
    parser = argparse.ArgumentParser(
        description='Deploy RocketGraph on AWS, Azure, or GCP using a profile file.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python deploy.py --profile config/default.yml
  python deploy.py --profile config/customer-a.yml --plan
  python deploy.py --profile config/customer-a.yml --destroy
  python deploy.py --list-profiles
        """
    )

    parser.add_argument('--profile', '-p',
                        default='config/default.yml',
                        help='Path to YAML or JSON deployment profile (default: config/default.yml)')
    parser.add_argument('--plan', action='store_true',
                        help='Run terraform plan only, do not apply')
    parser.add_argument('--destroy', '-d', action='store_true',
                        help='Destroy the deployed infrastructure')
    parser.add_argument('--auto-approve', '-y', action='store_true',
                        help='Skip interactive approval of terraform changes')
    parser.add_argument('--list-profiles', '-l', action='store_true',
                        help='List available profiles in config/ and exit')
    parser.add_argument('--dry-run', action='store_true',
                        help='Show secrets that would be pushed to the cloud store without '
                             'making any API calls or running Terraform')
    parser.add_argument('--secrets-only', action='store_true',
                        help='Push secrets to the cloud store and exit without running Terraform')
    parser.add_argument('--skip-secrets', action='store_true',
                        help='Skip pushing secrets to the cloud store (assume already done)')

    args = parser.parse_args()

    if args.list_profiles:
        list_profiles()
        return

    profile = load_profile(args.profile)

    # CLI flags override profile values
    if args.auto_approve:
        profile['auto_approve'] = True

    deployer = CloudDeployer(profile)

    try:
        deployer.validate_provider()
        deployer.check_prerequisites()
        deployer.validate_credentials()

        if args.dry_run:
            deployer.push_secrets(dry_run=True)
            return

        if not args.skip_secrets and not args.destroy:
            deployer.push_secrets(dry_run=False)

        if args.secrets_only:
            return

        deployer.generate_tfvars()

        if args.destroy:
            deployer.run_terraform('destroy')
        elif args.plan:
            deployer.run_terraform('plan')
        else:
            deployer.run_terraform('apply')

    except KeyboardInterrupt:
        print("\nDeployment cancelled")
        sys.exit(1)
    except Exception as e:
        print(f"\nError: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()