---
# AWS C2 Server Deployment Playbook

- name: Deploy AWS C2 server
  hosts: localhost
  gather_facts: false
  connection: local
  vars_files:
    - vars.yaml
  vars:
    # Default values
    ssh_user: "{{ ssh_user | default('kali') }}"
    aws_region: "{{ aws_region | default(aws_region_choices | random) }}"
    instance_type: "{{ aws_instance_type | default('t2.medium') }}"
    deployment_id: "{{ deployment_id | default('') }}"
    c2_name: "{{ c2_name | default('s-' + deployment_id) }}"
    # Define split_regions - only true when regions are explicitly different
    split_regions: "{{ c2_region is defined and redirector_region is defined and c2_region != redirector_region }}"
    # Only use shared infra when NOT doing split-region deployment
    use_shared_infra: "{{ not split_regions and not c2_only | default(false) | bool and not redirector_only | default(false) | bool }}"
    # Set correct region variable
    aws_c2_region: "{{ c2_region | default(aws_region) }}"
    # AMI map comes from vars.yaml - add fallback for safety
    kali_ami_map_fallback:
      us-east-1: "ami-061b17d332829ab1c"
      us-east-2: "ami-061b17d332829ab1c"  # Fallback to us-east-1 AMI

  tasks:
    - name: Validate AWS credentials
      assert:
        that:
          - aws_access_key is defined and aws_access_key != ""
          - aws_secret_key is defined and aws_secret_key != ""
        fail_msg: "AWS credentials are required"

    # Load shared infrastructure state if available
    - name: Check for shared infrastructure state
      stat:
        path: "infrastructure_state_{{ deployment_id }}.json"
      register: infra_state_file
      when: use_shared_infra | bool
      
    - name: Load shared infrastructure state
      include_vars:
        file: "infrastructure_state_{{ deployment_id }}.json"
        name: shared_infra
      when: use_shared_infra | bool and infra_state_file.stat.exists | default(false)
      
    - name: Set region for C2
      set_fact:
        aws_c2_region: "{{ shared_infra.region | default(aws_region) }}"
      when: use_shared_infra | bool and infra_state_file.stat.exists | default(false)
        
    - name: Set default region for C2
      set_fact:
        aws_c2_region: "{{ c2_region | default(aws_region) }}"
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    - name: Check if ami_map is provided in vars.yaml
      debug:
        msg: "ami_map is {{ 'defined' if ami_map is defined else 'NOT defined' }} in vars.yaml"

    - name: Set AMI ID for selected region (from vars.yaml)
      set_fact:
        ami_id: "{{ ami_map[aws_c2_region] | default(ami_map.us-east-1) }}"
      when: ami_map is defined and ami_map

    - name: Set AMI ID for selected region (fallback)
      set_fact:
        ami_id: "{{ kali_ami_map_fallback[aws_c2_region] | default(kali_ami_map_fallback['us-east-1']) }}"
      when: ami_id is not defined or ami_id == ""

    - name: Ensure we have a valid AMI ID
      assert:
        that:
          - ami_id is defined and ami_id != ""
        fail_msg: "Could not determine a valid AMI ID for region {{ aws_c2_region }}. Please add it to ami_map in vars.yaml."

    # Add AMI username mapping - improved with better detection
    - name: Determine correct SSH user for the AMI
      set_fact:
        ami_ssh_user: "{{ 'kali' if (ami_id is defined and ami_id is search('-kali-')) or (ami_id is defined and ami_id == 'ami-061b17d332829ab1c') else 'ubuntu' }}"

    - name: Display AMI and user information for debugging
      debug:
        msg: 
          - "Using AMI ID: {{ ami_id | default('AMI not defined') }}"
          - "Detected SSH user: {{ ami_ssh_user }}"

    # Create new infrastructure only if not using shared
    - name: Create VPC
      amazon.aws.ec2_vpc_net:
        name: "{{ c2_name }}-vpc"
        cidr_block: "10.0.0.0/16"
        region: "{{ aws_c2_region }}"
        tags:
          Name: "{{ c2_name }}-vpc"
          deployment_id: "{{ deployment_id }}"
        state: present
      register: vpc_result
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    - name: Create internet gateway for VPC
      amazon.aws.ec2_vpc_igw:
        vpc_id: "{{ vpc_result.vpc.id }}"
        region: "{{ aws_c2_region }}"
        state: present
        tags:
          Name: "{{ c2_name }}-igw"
          deployment_id: "{{ deployment_id }}"
      register: igw_result
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    - name: Create subnet in VPC
      amazon.aws.ec2_vpc_subnet:
        vpc_id: "{{ vpc_result.vpc.id }}"
        cidr: "10.0.1.0/24"
        region: "{{ aws_c2_region }}"
        az: "{{ aws_c2_region }}a"
        map_public: yes
        tags:
          Name: "{{ c2_name }}-subnet"
          deployment_id: "{{ deployment_id }}"
      register: subnet_result
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    - name: Create routing table for internet access
      amazon.aws.ec2_vpc_route_table:
        vpc_id: "{{ vpc_result.vpc.id }}"
        region: "{{ aws_c2_region }}"
        tags:
          Name: "{{ c2_name }}-rtb"
          deployment_id: "{{ deployment_id }}"
        routes:
          - dest: "0.0.0.0/0"
            gateway_id: "{{ igw_result.gateway_id }}"
        subnets:
          - "{{ subnet_result.subnet.id }}"
      register: route_table_result
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    # Set VPC ID based on shared or created
    - name: Set VPC ID from shared infrastructure
      set_fact:
        vpc_id: "{{ shared_infra.vpc_id }}"
        subnet_id: "{{ shared_infra.subnet_id }}"
        c2_vpc_id: "{{ shared_infra.vpc_id }}"  # Store for cleanup reference
      when: use_shared_infra | bool and infra_state_file.stat.exists | default(false)

    - name: Set VPC ID from created infrastructure
      set_fact:
        vpc_id: "{{ vpc_result.vpc.id }}"
        subnet_id: "{{ subnet_result.subnet.id }}"
        c2_vpc_id: "{{ vpc_result.vpc.id }}"  # Store for cleanup reference
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    - name: Set default redirector_ip for C2-only deployments
      set_fact:
        redirector_ip: "{{ operator_ip }}"
      when: c2_only | default(false) | bool and redirector_ip is not defined

    - name: Load redirector_ip from infrastructure state 
      block:
        - name: Check if infrastructure state file exists
          stat:
            path: "infrastructure_state_{{ deployment_id }}.json"
          register: infra_state_check
        
        - name: Load infrastructure state
          include_vars:
            file: "infrastructure_state_{{ deployment_id }}.json"
            name: infra_state
          when: infra_state_check.stat.exists
        
        - name: Set redirector IP from infrastructure state
          set_fact:
            redirector_ip: "{{ infra_state.redirector_ip | default(operator_ip) }}"
          when: infra_state_check.stat.exists and infra_state.redirector_ip is defined
        
        - name: Debug redirector IP
          debug:
            msg: "Using redirector IP: {{ redirector_ip | default('undefined') }}"
      when: redirector_ip is undefined

    - name: Default to operator IP if redirector IP is still undefined
      set_fact:
        redirector_ip: "{{ operator_ip }}"
      when: redirector_ip is undefined

    - name: Create security group for C2 server
      amazon.aws.ec2_security_group:
        name: "{{ c2_name }}-sg"
        description: "Secured C2 server {{ c2_name }}"
        vpc_id: "{{ vpc_id }}"
        region: "{{ aws_c2_region }}"
        rules:
          # Management access only from operator IP
          - proto: tcp
            ports: 22
            cidr_ip: "{{ operator_ip }}/32"
          - proto: tcp
            ports: "{{ havoc_teamserver_port | default(40056) }}"
            cidr_ip: "{{ operator_ip }}/32"
          # Allow traffic only from redirector
          - proto: tcp
            ports:
              - 80
              - 443
              - "{{ havoc_http_port | default(8080) }}"
              - "{{ havoc_https_port | default(443) }}"
              - "{{ gophish_admin_port }}"
            cidr_ip: "{{ redirector_ip }}/32"
        rules_egress:
          - proto: -1
            cidr_ip: 0.0.0.0/0
        state: present
      register: security_group

    # Generate or import SSH key for the deployment - FIXED KEY HANDLING
    - name: Check if deployment SSH key already exists locally
      stat:
        path: "~/.ssh/c2deploy_{{ deployment_id }}.pem"
      register: ssh_key_file

    - name: Generate key pair if it doesn't exist
      block:
        - name: Create SSH key pair
          command: ssh-keygen -t rsa -b 2048 -f ~/.ssh/c2deploy_{{ deployment_id }} -N ""
          args:
            creates: "~/.ssh/c2deploy_{{ deployment_id }}"
        
        - name: Rename private key to .pem format
          command: mv ~/.ssh/c2deploy_{{ deployment_id }} ~/.ssh/c2deploy_{{ deployment_id }}.pem
          args:
            creates: "~/.ssh/c2deploy_{{ deployment_id }}.pem" 
            removes: "~/.ssh/c2deploy_{{ deployment_id }}"
      when: not ssh_key_file.stat.exists

    - name: Ensure proper permissions on SSH key
      file:
        path: "~/.ssh/c2deploy_{{ deployment_id }}.pem"
        mode: '0600'
        state: file

    - name: Check if key pair exists in AWS
      amazon.aws.ec2_key_info:
        region: "{{ aws_c2_region }}"
        filters:
          key-name: "c2deploy_{{ deployment_id }}"
      register: existing_key_pair

    - name: Import SSH key to AWS
      amazon.aws.ec2_key:
        name: "c2deploy_{{ deployment_id }}"
        key_material: "{{ lookup('file', '~/.ssh/c2deploy_{{ deployment_id }}.pub') }}"
        region: "{{ aws_c2_region }}"
        state: present
      when: existing_key_pair.keypairs | length == 0

    # Launch the C2 server - use the consistent key pair name
    - name: Launch C2 instance
      amazon.aws.ec2_instance:
        name: "{{ c2_name }}"
        key_name: "c2deploy_{{ deployment_id }}"  # Use the imported key
        instance_type: "{{ instance_type | default('t2.medium') }}"
        vpc_subnet_id: "{{ subnet_id }}"
        security_groups:
          - "{{ security_group.group_id }}"
        image_id: "{{ ami_id }}"
        region: "{{ aws_c2_region }}"
        state: present
        wait: yes
        volumes:
          - device_name: "/dev/xvda"
            ebs:
              volume_size: 100
              delete_on_termination: true
        tags:
          Name: "{{ c2_name }}"
          deployment_id: "{{ deployment_id }}"
      register: c2_instance

    - name: Set c2_ip for later use
      set_fact:
        c2_ip: "{{ c2_instance.instances[0].public_ip_address }}"
        c2_instance_id: "{{ c2_instance.instances[0].instance_id }}"

    - name: Display C2 instance details for debugging
      debug:
        msg:
          - "C2 IP: {{ c2_ip }}"
          - "C2 Instance ID: {{ c2_instance_id }}"
          - "SSH User to use: {{ ami_ssh_user }}"
          - "SSH Key path: ~/.ssh/c2deploy_{{ deployment_id }}.pem"

    - name: Wait for C2 instance initialization
      pause:
        seconds: 180
      when: c2_instance.changed

    - name: Set correct permissions on SSH key
      file:
        path: "~/.ssh/c2deploy_{{ deployment_id }}.pem"
        mode: "0600"

    - name: Wait for C2 SSH to be available
      wait_for:
        host: "{{ c2_ip }}"
        port: 22
        delay: 30
        timeout: 300
        state: started

    - name: Test SSH and prepare remote environment
      block:
        - name: Ensure .ansible directory exists with proper permissions
          shell: |
            ssh -i ~/.ssh/c2deploy_{{ deployment_id }}.pem -o StrictHostKeyChecking=no {{ ami_ssh_user }}@{{ c2_ip }} "sudo mkdir -p /root/.ansible/tmp && sudo chmod 0700 /root/.ansible/tmp && sudo chown {{ ami_ssh_user }}:{{ ami_ssh_user }} /root/.ansible/tmp"
          register: ssh_prep
          until: ssh_prep is success
          retries: 5
          delay: 15
          ignore_errors: yes
          delegate_to: localhost
        
        - name: Display SSH preparation results
          debug:
            msg: "SSH preparation completed: {{ ssh_prep.stdout | default('No output') }}"

    # Test SSH connection directly to verify key is working
    - name: Test SSH connection to verify key
      shell: "ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i ~/.ssh/c2deploy_{{ deployment_id }}.pem {{ ami_ssh_user }}@{{ c2_ip }} 'echo SSH CONNECTION SUCCESSFUL'"
      register: ssh_test
      ignore_errors: yes

    - name: Display SSH test results
      debug:
        msg: "{{ ssh_test.stdout | default('SSH Connection failed!') }}"

    - name: Add C2 to inventory with updated SSH key path
      add_host:
        name: "c2"
        groups: "c2servers"
        ansible_host: "{{ c2_ip }}"
        ansible_user: "{{ ami_ssh_user }}"
        ansible_ssh_private_key_file: "~/.ssh/c2deploy_{{ deployment_id }}.pem"
        ansible_ssh_common_args: "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes"
        ansible_python_interpreter: "/usr/bin/python3"
        # Add these lines to pass all required variables:
        smtp_auth_user: "{{ smtp_auth_user }}"
        smtp_auth_pass: "{{ smtp_auth_pass }}"
        gophish_admin_port: "{{ gophish_admin_port }}"
        domain: "{{ domain }}"
        redirector_subdomain: "{{ redirector_subdomain }}"
        letsencrypt_email: "{{ letsencrypt_email }}"
        havoc_teamserver_port: "{{ havoc_teamserver_port | default(40056) }}"
        havoc_http_port: "{{ havoc_http_port | default(8080) }}"
        havoc_https_port: "{{ havoc_https_port | default(443) }}"
        zero_logs: "{{ zero_logs | default(true) }}"
        secure_memory: "{{ secure_memory | default(true) }}"
        disable_history: "{{ disable_history | default(true) }}"
        setup_integrated_tracker: "{{ setup_integrated_tracker | default(false) }}"
        tracker_domain: "{{ tracker_domain | default('track.' + domain) | default('') }}"

# Configure C2 server with a proper structure
- name: Configure C2 server
  hosts: c2servers
  become: yes
  become_user: root
  gather_facts: true
  vars_files:
    - vars.yaml  # Add this line to load the variables
  vars:
    redirector_ip: "{{ hostvars['localhost']['redirector_ip'] | default('127.0.0.1') }}"
    c2_subdomain: "{{ c2_subdomain | default('mail') }}"
  tasks:
    - name: Install python3 if it doesn't exist on target
      raw: test -e /usr/bin/python3 || (apt-get update && apt-get install -y python3)
      args:
        executable: /bin/bash
      register: python_install
      ignore_errors: yes
      
    - name: Debug connection information
      debug:
        msg:
          - "Connected to C2 server successfully"
          - "Host: {{ ansible_host }}"
          - "User: {{ ansible_user }}"
          - "Python version: {{ ansible_python_version | default('unknown') }}"
      
    - name: Download Kali archive keyring to temporary location
      get_url:
        url: https://archive.kali.org/archive-keyring.gpg
        dest: /tmp/kali-archive-keyring.gpg
        mode: "0644"
        force: yes
      register: keyring_download

    - name: Install Kali archive keyring
      command: install -m 0644 /tmp/kali-archive-keyring.gpg /usr/share/keyrings/kali-archive-keyring.gpg
      when: keyring_download is changed

    - name: Wait for apt to be available
      apt:
        update_cache: yes
      register: apt_result
      until: apt_result is success
      retries: 10
      delay: 10

    - name: Install Rust compiler
      apt:
        name:
          - cargo
          - rustc
          - libssl-dev
          - pkg-config
        state: present
      register: rust_install
      until: rust_install is success
      retries: 3
      delay: 5

    - name: Clean up any failed pipx installations
      file:
        path: "{{ item }}"
        state: absent
      with_items:
        - "/root/.local/state/pipx/venvs/netexec"
        - "/root/.local/state/pipx/venvs/trevorspray"
      ignore_errors: yes

    - name: Include common tool installation tasks
      include_tasks: "../tasks/install_tools.yml"

    - name: Include common C2 configuration tasks
      include_tasks: "../tasks/configure_c2.yml"

    - name: Include common security hardening tasks
      include_tasks: "../tasks/security_hardening.yml"

    - name: Include common mail server configuration tasks
      include_tasks: "../tasks/configure_mail.yml"

    - name: Set up SSH access to redirector
      block:
        - name: Ensure /root/.ssh directory exists on C2 server
          file:
            path: /root/.ssh
            state: directory
            mode: '0700'
            owner: root
            group: root

        - name: Copy private SSH key to C2 server for redirector access
          copy:
            src: "{{ playbook_dir }}/../.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
            dest: "/root/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
            mode: '0600'
            owner: root
            group: root
          register: key_copy
          ignore_errors: yes

        - name: If direct path fails, try home directory location
          copy:
            src: "~/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
            dest: "/root/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
            mode: '0600'
            owner: root
            group: root
          when: key_copy is failed

        - name: Create SSH config file to use key automatically
          copy:
            dest: "/root/.ssh/config"
            content: |
              Host redirector
                HostName {{ hostvars['localhost']['redirector_ip'] }}
                User ubuntu
                IdentityFile /root/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem
                StrictHostKeyChecking no
                UserKnownHostsFile /dev/null
            mode: '0600'
            owner: root
            group: root

        - name: Create alias for easy redirector access
          lineinfile:
            path: /root/.bashrc
            line: 'alias redirector="ssh -i /root/.ssh/c2deploy_{{ hostvars["localhost"]["deployment_id"] }}.pem ubuntu@{{ hostvars["localhost"]["redirector_ip"] }}"'
            state: present

        - name: Test SSH from C2 to redirector
          shell: |
            ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i /root/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem ubuntu@{{ hostvars['localhost']['redirector_ip'] }} "echo 'SSH CONNECTION SUCCESSFUL FROM C2'"
          register: ssh_test_result
          changed_when: false
          ignore_errors: yes

        - name: Display SSH test result
          debug:
            msg: "{{ ssh_test_result.stdout if ssh_test_result.rc == 0 else 'SSH connection failed: ' + ssh_test_result.stderr }}"

    - name: Print deployment summary
      debug:
        msg:
          - "C2 Server Deployment Complete!"
          - "-----------------------------"
          - "C2 Server IP: {{ ansible_host }}"
          - "C2 Server Domain: {{ c2_subdomain }}.{{ domain }} (Update DNS A record)"
          - "GoPhish Admin Port: {{ gophish_admin_port }}"
          - "SSH Key: ~/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
      when: not disable_summary | default(false)