---
# AWS Redirector Deployment Playbook

- name: Deploy AWS redirector
  hosts: localhost
  gather_facts: false
  connection: local
  vars_files:
    - vars.yaml
  vars:
    # Default values
    ssh_user: "{{ ssh_user | default('ubuntu') }}"
    aws_region: "{{ aws_region | default(aws_region_choices | random) }}"
    instance_type: "{{ aws_instance_type | default('t2.micro') }}"
    deployment_id: "{{ deployment_id | default('') }}"
    redirector_name: "{{ redirector_name | default('r-' + deployment_id) }}"
    effective_listen_port: "{{ shell_handler_port | default(4488) }}"
    # Check for shared infrastructure
    # Only use shared when C2 and redirector are in the same region
    use_shared_infra: "{{ not (c2_region is defined and redirector_region is defined and c2_region != redirector_region) and not c2_only | default(false) | bool and not redirector_only | default(false) | bool }}"
    # Ubuntu AMI IDs for different regions (Ubuntu 22.04 LTS)
    ubuntu_ami_map:
      us-east-1: "ami-0aa2b7722dc1b5612"
      us-east-2: "ami-06c4532923d4ba1ec"
      us-west-1: "ami-0573b70afecda915d"
      us-west-2: "ami-0c79c59ac2c572b87"
      eu-west-1: "ami-0694d931cee176e7d"
      eu-west-2: "ami-0505148b3591e4c07"
      eu-central-1: "ami-06dd92ecc74fdfb36"
      ap-southeast-1: "ami-0df7a207adb9748c7"
      ap-southeast-2: "ami-0df4b2961410d4cff"
      ap-northeast-1: "ami-0014b5f031a76c1b1"
      sa-east-1: "ami-0af6e9042ea5a4e3e"

  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"

    # Set region for redirector - fix for noop task error
    - name: Set region for redirector
      set_fact:
        aws_redirector_region: "{{ redirector_region | default(aws_region) }}"

    # Load shared infrastructure state if available - fixed implementation
    - name: Check for shared infrastructure state
      block:
        - name: Check if state file exists
          stat:
            path: "infrastructure_state_{{ deployment_id }}.json"
          register: infra_state_file
        
        - name: Include vars if file exists
          include_vars:
            file: "infrastructure_state_{{ deployment_id }}.json"
            name: shared_infra
          when: infra_state_file.stat.exists | default(false)
      when: use_shared_infra | bool

    # After loading shared infrastructure state
    - name: Validate shared VPC exists
      amazon.aws.ec2_vpc_net_info:
        region: "{{ shared_infra.region }}"
        vpc_ids:
          - "{{ shared_infra.vpc_id }}"
      register: vpc_check
      when: use_shared_infra | bool and infra_state_file.stat.exists | default(false)
      ignore_errors: yes

    - name: Delete stale infrastructure state file
      file:
        path: "infrastructure_state.json"
        state: absent
      when: use_shared_infra | bool and vpc_check.vpcs is defined and vpc_check.vpcs | length == 0

    - name: Disable shared infrastructure when VPC doesn't exist
      set_fact:
        use_shared_infra: false
      when: use_shared_infra | bool and vpc_check.vpcs is defined and vpc_check.vpcs | length == 0

    - name: Set region variables from shared infra
      set_fact:
        aws_redirector_region: "{{ shared_infra.region | default(aws_region) }}"
      when: use_shared_infra | bool and infra_state_file.stat.exists | default(false)
        
    - name: Set AMI ID for selected region
      set_fact:
        ami_id: "{{ ubuntu_ami_map[aws_redirector_region] | default(ubuntu_ami_map['us-east-1']) }}"
      
    # Create new infrastructure only if not using shared
    - name: Create VPC
      amazon.aws.ec2_vpc_net:
        name: "{{ redirector_name }}-vpc"
        cidr_block: "10.0.0.0/16"
        region: "{{ aws_redirector_region }}"
        tags:
          Name: "{{ redirector_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_redirector_region }}"
        state: present
        tags:
          Name: "{{ redirector_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_redirector_region }}"
        az: "{{ aws_redirector_region }}a"
        map_public: yes
        tags:
          Name: "{{ redirector_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_redirector_region }}"
        tags:
          Name: "{{ redirector_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 }}"
        redirector_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 }}"
        redirector_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 c2_ip for redirector-only deployments
      set_fact:
        c2_ip: "{{ operator_ip }}"
      when: redirector_only | default(false) | bool and c2_ip is not defined

    - name: Load c2_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 c2 IP from infrastructure state
          set_fact:
            c2_ip: "{{ infra_state.c2_ip | default(operator_ip) }}"
          when: infra_state_check.stat.exists and infra_state.c2_ip is defined
        
        - name: Debug c2 IP
          debug:
            msg: "Using c2 IP: {{ c2_ip | default('undefined') }}"
      when: c2_ip is undefined

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

    - name: Create security group for redirector
      amazon.aws.ec2_security_group:
        name: "{{ redirector_name }}-sg"
        description: "Security group for redirector {{ redirector_name }}"
        vpc_id: "{{ vpc_id }}"
        region: "{{ aws_redirector_region }}"
        rules:
          # Management access only from operator IP
          - proto: tcp
            ports: 22
            cidr_ip: "{{ operator_ip }}/32"
          # Public-facing services are open to internet
          - proto: tcp
            ports:
              - 80
              - 443
              - "{{ effective_listen_port }}"
            cidr_ip: 0.0.0.0/0
        rules_egress:
          - proto: -1
            cidr_ip: 0.0.0.0/0
        state: present
      register: security_group

    # Save infrastructure state for reuse
    - name: Save infrastructure state for reuse
      copy:
        content: |
          {
            "vpc_id": "{{ vpc_id }}",
            "subnet_id": "{{ subnet_id }}",
            "security_group_id": "{{ security_group.group_id }}",
            "region": "{{ aws_redirector_region }}",
            "deployment_id": "{{ deployment_id }}"
          }
        dest: "infrastructure_state_{{ deployment_id }}.json"
      when: not (use_shared_infra | bool and infra_state_file.stat.exists | default(false))

    # 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 SSH 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
      
    # Check if key pair exists in AWS
    - name: Check if key pair exists in AWS
      amazon.aws.ec2_key_info:
        region: "{{ aws_redirector_region }}"
        filters:
          key-name: "c2deploy_{{ deployment_id }}"
      register: existing_key_pair

    # Import SSH key to AWS - FIXED CONDITION
    - 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_redirector_region }}"
        state: present
      when: existing_key_pair.keypairs | length == 0

    # Launch the redirector instance
    - name: Launch redirector instance
      amazon.aws.ec2_instance:
        name: "{{ redirector_name }}"
        key_name: "c2deploy_{{ deployment_id }}"
        instance_type: "{{ instance_type | default('t2.micro') }}"
        vpc_subnet_id: "{{ subnet_id }}"
        security_groups:
          - "{{ security_group.group_id }}"
        image_id: "{{ ami_id }}"
        region: "{{ aws_redirector_region }}"
        state: present
        wait: yes
        volumes:
          - device_name: "/dev/xvda"
            ebs:
              volume_size: 20
              delete_on_termination: true
        tags:
          Name: "{{ redirector_name }}"
          deployment_id: "{{ deployment_id }}"
      register: redirector_instance

    - name: Set redirector_ip for later use
      set_fact:
        redirector_ip: "{{ redirector_instance.instances[0].public_ip_address }}"
        redirector_instance_id: "{{ redirector_instance.instances[0].instance_id }}"

    - name: Display redirector instance details for debugging
      debug:
        msg:
          - "Redirector IP: {{ redirector_ip }}"
          - "Redirector Instance ID: {{ redirector_instance_id }}"
          - "SSH User to use: {{ ssh_user }}"
          - "SSH Key path: ~/.ssh/c2deploy_{{ deployment_id }}.pem"

    - name: Wait for instance initialization
      pause:
        seconds: 120
      when: redirector_instance.changed

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

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

    - name: Add redirector to inventory
      add_host:
        name: "redirector"
        groups: "redirectors"
        ansible_host: "{{ redirector_ip }}"
        ansible_user: "ubuntu"
        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"
        shell_handler_port: "{{ effective_listen_port }}"
        # Add necessary variables for redirector configuration
        c2_ip: "{{ c2_ip }}"
        domain: "{{ domain }}"
        redirector_subdomain: "{{ redirector_subdomain }}"
        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) }}"

# Rest of the playbook for configuring the redirector
- name: Configure redirector
  hosts: redirectors
  become: true
  gather_facts: true
  vars_files:
    - vars.yaml
  vars:
    c2_ip: "{{ hostvars['localhost']['c2_ip'] | default('127.0.0.1') }}"
    shell_handler_port: "{{ hostvars['localhost']['effective_listen_port'] }}"
    
  # Include the rest of your redirector configuration tasks here
  tasks:    
    - name: Include common redirector configuration tasks
      include_tasks: "../../modules/redirectors/tasks/configure_redirector.yml"
      
    - name: Configure shell handler script with listening port
      template:
        src: "../../modules/c2/files/havoc_shell_handler.sh"
        dest: "/root/Tools/shell_handler.sh"
        mode: 0755
      vars:
        listen_port: "{{ effective_listen_port }}"
      
    - name: Print deployment summary
      debug:
        msg:
          - "Redirector Deployment Complete!"
          - "-----------------------------"
          - "Redirector IP: {{ ansible_host }}"
          - "Redirector Domain: {{ redirector_subdomain | default('cdn') }}.{{ domain }} (Update DNS A record)"
          - "SSH Key: ~/.ssh/c2deploy_{{ hostvars['localhost']['deployment_id'] }}.pem"
      when: not disable_summary | default(false)