# aws.py - Automation code for Amazon Web Services

#

# Copyright 2019 Gregory Szorc <gregory.szorc@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

# no-check-code because Python 3 native.

import contextlib

import copy

import hashlib

import json

import os

import pathlib

import subprocess

import time

import boto3

import botocore.exceptions

from .linux import BOOTSTRAP_DEBIAN

from .ssh import (

    exec_command as ssh_exec_command,

    wait_for_ssh,

)

from .winrm import (

    run_powershell,

    wait_for_winrm,

)

SOURCE_ROOT = pathlib.Path(

    os.path.abspath(__file__)

).parent.parent.parent.parent

INSTALL_WINDOWS_DEPENDENCIES = (

    SOURCE_ROOT / 'contrib' / 'install-windows-dependencies.ps1'

)

INSTANCE_TYPES_WITH_STORAGE = {

    'c5d',

    'd2',

    'h1',

    'i3',

    'm5ad',

    'm5d',

    'r5d',

    'r5ad',

    'x1',

    'z1d',

}

AMAZON_ACCOUNT_ID = '801119661308'

DEBIAN_ACCOUNT_ID = '379101102735'

DEBIAN_ACCOUNT_ID_2 = '136693071363'

UBUNTU_ACCOUNT_ID = '099720109477'

WINDOWS_BASE_IMAGE_NAME = 'Windows_Server-2022-English-Full-Base-*'

KEY_PAIRS = {

    'automation',

}

SECURITY_GROUPS = {

    'linux-dev-1': {

        'description': 'Mercurial Linux instances that perform build/test automation',

        'ingress': [

            {

                'FromPort': 22,

                'ToPort': 22,

                'IpProtocol': 'tcp',

                'IpRanges': [

                    {

                        'CidrIp': '0.0.0.0/0',

                        'Description': 'SSH from entire Internet',

                    },

                ],

            },

        ],

    },

    'windows-dev-1': {

        'description': 'Mercurial Windows instances that perform build automation',

        'ingress': [

            {

                'FromPort': 22,

                'ToPort': 22,

                'IpProtocol': 'tcp',

                'IpRanges': [

                    {

                        'CidrIp': '0.0.0.0/0',

                        'Description': 'SSH from entire Internet',

                    },

                ],

            },

            {

                'FromPort': 3389,

                'ToPort': 3389,

                'IpProtocol': 'tcp',

                'IpRanges': [

                    {

                        'CidrIp': '0.0.0.0/0',

                        'Description': 'RDP from entire Internet',

                    },

                ],

            },

            {

                'FromPort': 5985,

                'ToPort': 5986,

                'IpProtocol': 'tcp',

                'IpRanges': [

                    {

                        'CidrIp': '0.0.0.0/0',

                        'Description': 'PowerShell Remoting (Windows Remote Management)',

                    },

                ],

            },

        ],

    },

}

IAM_ROLES = {

    'ephemeral-ec2-role-1': {

        'description': 'Mercurial temporary EC2 instances',

        'policy_arns': [

            'arn:aws:iam::aws:policy/service-role/AmazonEC2RoleforSSM',

        ],

    },

}

ASSUME_ROLE_POLICY_DOCUMENT = '''

{

  "Version": "2012-10-17",

  "Statement": [

    {

      "Effect": "Allow",

      "Principal": {

        "Service": "ec2.amazonaws.com"

      },

      "Action": "sts:AssumeRole"

    }

  ]

}

'''.strip()

IAM_INSTANCE_PROFILES = {

    'ephemeral-ec2-1': {

        'roles': [

            'ephemeral-ec2-role-1',

        ],

    }

}

# User Data for Windows EC2 instance. Mainly used to set the password

# and configure WinRM.

# Inspired by the User Data script used by Packer

# (from https://www.packer.io/intro/getting-started/build-image.html).

WINDOWS_USER_DATA = r'''

<powershell>

# TODO enable this once we figure out what is failing.

#$ErrorActionPreference = "stop"

# Set administrator password

net user Administrator "%s"

wmic useraccount where "name='Administrator'" set PasswordExpires=FALSE

# And set it via EC2Launch so it persists across reboots.

$config = & $env:ProgramFiles\Amazon\EC2Launch\EC2Launch.exe get-agent-config --format json | ConvertFrom-Json

$config | ConvertTo-Json -Depth 6 | Out-File -encoding UTF8 $env:ProgramData/Amazon/EC2Launch/config/agent-config.yml

$setAdminAccount = @"

{

  "task": "setAdminAccount",

  "inputs": {

    "password": {

      "type": "static",

      "data": "%s"

    }

  }

}

"@

$config.config | %%{if($_.stage -eq 'preReady'){$_.tasks += (ConvertFrom-Json -InputObject $setAdminAccount)}}

$config | ConvertTo-Json -Depth 6 | Out-File -encoding UTF8 $env:ProgramData/Amazon/EC2Launch/config/agent-config.yml

# First, make sure WinRM can't be connected to

netsh advfirewall firewall set rule name="Windows Remote Management (HTTP-In)" new enable=yes action=block

# Delete any existing WinRM listeners

winrm delete winrm/config/listener?Address=*+Transport=HTTP  2>$Null

winrm delete winrm/config/listener?Address=*+Transport=HTTPS 2>$Null

# Create a new WinRM listener and configure

winrm create winrm/config/listener?Address=*+Transport=HTTP

winrm set winrm/config/winrs '@{MaxMemoryPerShellMB="0"}'

winrm set winrm/config '@{MaxTimeoutms="7200000"}'

winrm set winrm/config/service '@{AllowUnencrypted="true"}'

winrm set winrm/config/service '@{MaxConcurrentOperationsPerUser="12000"}'

winrm set winrm/config/service/auth '@{Basic="true"}'

winrm set winrm/config/client/auth '@{Basic="true"}'

# Configure UAC to allow privilege elevation in remote shells

$Key = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System'

$Setting = 'LocalAccountTokenFilterPolicy'

Set-ItemProperty -Path $Key -Name $Setting -Value 1 -Force

# Avoid long usernames in the temp directory path because the '~' causes extra quoting in ssh output

[System.Environment]::SetEnvironmentVariable('TMP', 'C:\Temp', [System.EnvironmentVariableTarget]::User)

[System.Environment]::SetEnvironmentVariable('TEMP', 'C:\Temp', [System.EnvironmentVariableTarget]::User)

# Configure and restart the WinRM Service; Enable the required firewall exception

Stop-Service -Name WinRM

Set-Service -Name WinRM -StartupType Automatic

netsh advfirewall firewall set rule name="Windows Remote Management (HTTP-In)" new action=allow localip=any remoteip=any

Start-Service -Name WinRM

# Disable firewall on private network interfaces so prompts don't appear.

Set-NetFirewallProfile -Name private -Enabled false

</powershell>

'''.lstrip()

WINDOWS_BOOTSTRAP_POWERSHELL = '''

Write-Output "installing PowerShell dependencies"

Install-PackageProvider -Name NuGet -MinimumVersion 2.8.5.201 -Force

Set-PSRepository -Name PSGallery -InstallationPolicy Trusted

Install-Module -Name OpenSSHUtils -RequiredVersion 0.0.2.0

Write-Output "installing OpenSSL server"

Add-WindowsCapability -Online -Name OpenSSH.Server~~~~0.0.1.0

# Various tools will attempt to use older versions of .NET. So we enable

# the feature that provides them so it doesn't have to be auto-enabled

# later.

Write-Output "enabling .NET Framework feature"

Install-WindowsFeature -Name Net-Framework-Core

'''

class AWSConnection:

    """Manages the state of a connection with AWS."""

    def __init__(self, automation, region: str, ensure_ec2_state: bool = True):

        self.automation = automation

        self.local_state_path = automation.state_path

        self.prefix = 'hg-'

        self.session = boto3.session.Session(region_name=region)

        self.ec2client = self.session.client('ec2')

        self.ec2resource = self.session.resource('ec2')

        self.iamclient = self.session.client('iam')

        self.iamresource = self.session.resource('iam')

        self.security_groups = {}

        if ensure_ec2_state:

            ensure_key_pairs(automation.state_path, self.ec2resource)

            self.security_groups = ensure_security_groups(self.ec2resource)

            ensure_iam_state(self.iamclient, self.iamresource)

    def key_pair_path_private(self, name):

        """Path to a key pair private key file."""

        return self.local_state_path / 'keys' / ('keypair-%s' % name)

    def key_pair_path_public(self, name):

        return self.local_state_path / 'keys' / ('keypair-%s.pub' % name)

def rsa_key_fingerprint(p: pathlib.Path):

    """Compute the fingerprint of an RSA private key."""

    # TODO use rsa package.

    res = subprocess.run(

        [

            'openssl',

            'pkcs8',

            '-in',

            str(p),

            '-nocrypt',

            '-topk8',

            '-outform',

            'DER',

        ],

        capture_output=True,

        check=True,

    )

    sha1 = hashlib.sha1(res.stdout).hexdigest()

    return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2]))

def ensure_key_pairs(state_path: pathlib.Path, ec2resource, prefix='hg-'):

    remote_existing = {}

    for kpi in ec2resource.key_pairs.all():

        if kpi.name.startswith(prefix):

            remote_existing[kpi.name[len(prefix) :]] = kpi.key_fingerprint

    # Validate that we have these keys locally.

    key_path = state_path / 'keys'

    key_path.mkdir(exist_ok=True, mode=0o700)

    def remove_remote(name):

        print('deleting key pair %s' % name)

        key = ec2resource.KeyPair(name)

        key.delete()

    def remove_local(name):

        pub_full = key_path / ('keypair-%s.pub' % name)

        priv_full = key_path / ('keypair-%s' % name)

        print('removing %s' % pub_full)

        pub_full.unlink()

        print('removing %s' % priv_full)

        priv_full.unlink()

    local_existing = {}

    for f in sorted(os.listdir(key_path)):

        if not f.startswith('keypair-') or not f.endswith('.pub'):

            continue

        name = f[len('keypair-') : -len('.pub')]

        pub_full = key_path / f

        priv_full = key_path / ('keypair-%s' % name)

        with open(pub_full, 'r', encoding='ascii') as fh:

            data = fh.read()

        if not data.startswith('ssh-rsa '):

            print(

                'unexpected format for key pair file: %s; removing' % pub_full

            )

            pub_full.unlink()

            priv_full.unlink()

            continue

        local_existing[name] = rsa_key_fingerprint(priv_full)

    for name in sorted(set(remote_existing) | set(local_existing)):

        if name not in local_existing:

            actual = '%s%s' % (prefix, name)

            print('remote key %s does not exist locally' % name)

            remove_remote(actual)

            del remote_existing[name]

        elif name not in remote_existing:

            print('local key %s does not exist remotely' % name)

            remove_local(name)

            del local_existing[name]

        elif remote_existing[name] != local_existing[name]:

            print(

                'key fingerprint mismatch for %s; '

                'removing from local and remote' % name

            )

            remove_local(name)

            remove_remote('%s%s' % (prefix, name))

            del local_existing[name]

            del remote_existing[name]

    missing = KEY_PAIRS - set(remote_existing)

    for name in sorted(missing):

        actual = '%s%s' % (prefix, name)

        print('creating key pair %s' % actual)

        priv_full = key_path / ('keypair-%s' % name)

        pub_full = key_path / ('keypair-%s.pub' % name)

        kp = ec2resource.create_key_pair(KeyName=actual)

        with priv_full.open('w', encoding='ascii') as fh:

            fh.write(kp.key_material)

            fh.write('\n')

        priv_full.chmod(0o0600)

        # SSH public key can be extracted via `ssh-keygen`.

        with pub_full.open('w', encoding='ascii') as fh:

            subprocess.run(

                ['ssh-keygen', '-y', '-f', str(priv_full)],

                stdout=fh,

                check=True,

            )

        pub_full.chmod(0o0600)

def delete_instance_profile(profile):

    for role in profile.roles:

        print(

            'removing role %s from instance profile %s'

            % (role.name, profile.name)

        )

        profile.remove_role(RoleName=role.name)

    print('deleting instance profile %s' % profile.name)

    profile.delete()

def ensure_iam_state(iamclient, iamresource, prefix='hg-'):

    """Ensure IAM state is in sync with our canonical definition."""

    remote_profiles = {}

    for profile in iamresource.instance_profiles.all():

        if profile.name.startswith(prefix):

            remote_profiles[profile.name[len(prefix) :]] = profile

    for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)):

        delete_instance_profile(remote_profiles[name])

        del remote_profiles[name]

    remote_roles = {}

    for role in iamresource.roles.all():

        if role.name.startswith(prefix):

            remote_roles[role.name[len(prefix) :]] = role