diff contrib/automation/hgautomation/aws.py @ 43076:2372284d9457

formatting: blacken the codebase This is using my patch to black (https://github.com/psf/black/pull/826) so we don't un-wrap collection literals. Done with: hg files 'set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"' | xargs black -S # skip-blame mass-reformatting only # no-check-commit reformats foo_bar functions Differential Revision: https://phab.mercurial-scm.org/D6971
author Augie Fackler <augie@google.com>
date Sun, 06 Oct 2019 09:45:02 -0400
parents d1d919f679f7
children c09e8ac3f61f
line wrap: on
line diff
--- a/contrib/automation/hgautomation/aws.py	Sat Oct 05 10:29:34 2019 -0400
+++ b/contrib/automation/hgautomation/aws.py	Sun Oct 06 09:45:02 2019 -0400
@@ -19,9 +19,7 @@
 import boto3
 import botocore.exceptions
 
-from .linux import (
-    BOOTSTRAP_DEBIAN,
-)
+from .linux import BOOTSTRAP_DEBIAN
 from .ssh import (
     exec_command as ssh_exec_command,
     wait_for_ssh,
@@ -32,10 +30,13 @@
 )
 
 
-SOURCE_ROOT = pathlib.Path(os.path.abspath(__file__)).parent.parent.parent.parent
+SOURCE_ROOT = pathlib.Path(
+    os.path.abspath(__file__)
+).parent.parent.parent.parent
 
-INSTALL_WINDOWS_DEPENDENCIES = (SOURCE_ROOT / 'contrib' /
-                                'install-windows-dependencies.ps1')
+INSTALL_WINDOWS_DEPENDENCIES = (
+    SOURCE_ROOT / 'contrib' / 'install-windows-dependencies.ps1'
+)
 
 
 INSTANCE_TYPES_WITH_STORAGE = {
@@ -107,7 +108,6 @@
                         'Description': 'RDP from entire Internet',
                     },
                 ],
-
             },
             {
                 'FromPort': 5985,
@@ -119,7 +119,7 @@
                         'Description': 'PowerShell Remoting (Windows Remote Management)',
                     },
                 ],
-            }
+            },
         ],
     },
 }
@@ -152,11 +152,7 @@
 
 
 IAM_INSTANCE_PROFILES = {
-    'ephemeral-ec2-1': {
-        'roles': [
-            'ephemeral-ec2-role-1',
-        ],
-    }
+    'ephemeral-ec2-1': {'roles': ['ephemeral-ec2-role-1',],}
 }
 
 
@@ -226,7 +222,7 @@
 class AWSConnection:
     """Manages the state of a connection with AWS."""
 
-    def __init__(self, automation, region: str, ensure_ec2_state: bool=True):
+    def __init__(self, automation, region: str, ensure_ec2_state: bool = True):
         self.automation = automation
         self.local_state_path = automation.state_path
 
@@ -257,10 +253,19 @@
 
     # TODO use rsa package.
     res = subprocess.run(
-        ['openssl', 'pkcs8', '-in', str(p), '-nocrypt', '-topk8',
-         '-outform', 'DER'],
+        [
+            'openssl',
+            'pkcs8',
+            '-in',
+            str(p),
+            '-nocrypt',
+            '-topk8',
+            '-outform',
+            'DER',
+        ],
         capture_output=True,
-        check=True)
+        check=True,
+    )
 
     sha1 = hashlib.sha1(res.stdout).hexdigest()
     return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2]))
@@ -271,7 +276,7 @@
 
     for kpi in ec2resource.key_pairs.all():
         if kpi.name.startswith(prefix):
-            remote_existing[kpi.name[len(prefix):]] = kpi.key_fingerprint
+            remote_existing[kpi.name[len(prefix) :]] = kpi.key_fingerprint
 
     # Validate that we have these keys locally.
     key_path = state_path / 'keys'
@@ -297,7 +302,7 @@
         if not f.startswith('keypair-') or not f.endswith('.pub'):
             continue
 
-        name = f[len('keypair-'):-len('.pub')]
+        name = f[len('keypair-') : -len('.pub')]
 
         pub_full = key_path / f
         priv_full = key_path / ('keypair-%s' % name)
@@ -306,8 +311,9 @@
             data = fh.read()
 
         if not data.startswith('ssh-rsa '):
-            print('unexpected format for key pair file: %s; removing' %
-                  pub_full)
+            print(
+                'unexpected format for key pair file: %s; removing' % pub_full
+            )
             pub_full.unlink()
             priv_full.unlink()
             continue
@@ -327,8 +333,10 @@
             del local_existing[name]
 
         elif remote_existing[name] != local_existing[name]:
-            print('key fingerprint mismatch for %s; '
-                  'removing from local and remote' % 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]
@@ -356,15 +364,18 @@
             subprocess.run(
                 ['ssh-keygen', '-y', '-f', str(priv_full)],
                 stdout=fh,
-                check=True)
+                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))
+        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)
@@ -378,7 +389,7 @@
 
     for profile in iamresource.instance_profiles.all():
         if profile.name.startswith(prefix):
-            remote_profiles[profile.name[len(prefix):]] = profile
+            remote_profiles[profile.name[len(prefix) :]] = profile
 
     for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)):
         delete_instance_profile(remote_profiles[name])
@@ -388,7 +399,7 @@
 
     for role in iamresource.roles.all():
         if role.name.startswith(prefix):
-            remote_roles[role.name[len(prefix):]] = role
+            remote_roles[role.name[len(prefix) :]] = role
 
     for name in sorted(set(remote_roles) - set(IAM_ROLES)):
         role = remote_roles[name]
@@ -404,7 +415,8 @@
         print('creating IAM instance profile %s' % actual)
 
         profile = iamresource.create_instance_profile(
-            InstanceProfileName=actual)
+            InstanceProfileName=actual
+        )
         remote_profiles[name] = profile
 
         waiter = iamclient.get_waiter('instance_profile_exists')
@@ -453,23 +465,12 @@
 
     images = ec2resource.images.filter(
         Filters=[
-            {
-                'Name': 'owner-id',
-                'Values': [owner_id],
-            },
-            {
-                'Name': 'state',
-                'Values': ['available'],
-            },
-            {
-                'Name': 'image-type',
-                'Values': ['machine'],
-            },
-            {
-                'Name': 'name',
-                'Values': [name],
-            },
-        ])
+            {'Name': 'owner-id', 'Values': [owner_id],},
+            {'Name': 'state', 'Values': ['available'],},
+            {'Name': 'image-type', 'Values': ['machine'],},
+            {'Name': 'name', 'Values': [name],},
+        ]
+    )
 
     for image in images:
         return image
@@ -487,7 +488,7 @@
 
     for group in ec2resource.security_groups.all():
         if group.group_name.startswith(prefix):
-            existing[group.group_name[len(prefix):]] = group
+            existing[group.group_name[len(prefix) :]] = group
 
     purge = set(existing) - set(SECURITY_GROUPS)
 
@@ -507,13 +508,10 @@
         print('adding security group %s' % actual)
 
         group_res = ec2resource.create_security_group(
-            Description=group['description'],
-            GroupName=actual,
+            Description=group['description'], GroupName=actual,
         )
 
-        group_res.authorize_ingress(
-            IpPermissions=group['ingress'],
-        )
+        group_res.authorize_ingress(IpPermissions=group['ingress'],)
 
         security_groups[name] = group_res
 
@@ -577,8 +575,10 @@
                 instance.reload()
                 continue
 
-            print('public IP address for %s: %s' % (
-                instance.id, instance.public_ip_address))
+            print(
+                'public IP address for %s: %s'
+                % (instance.id, instance.public_ip_address)
+            )
             break
 
 
@@ -603,10 +603,7 @@
     while True:
         res = ssmclient.describe_instance_information(
             Filters=[
-                {
-                    'Key': 'InstanceIds',
-                    'Values': [i.id for i in instances],
-                },
+                {'Key': 'InstanceIds', 'Values': [i.id for i in instances],},
             ],
         )
 
@@ -628,9 +625,7 @@
         InstanceIds=[i.id for i in instances],
         DocumentName=document_name,
         Parameters=parameters,
-        CloudWatchOutputConfig={
-            'CloudWatchOutputEnabled': True,
-        },
+        CloudWatchOutputConfig={'CloudWatchOutputEnabled': True,},
     )
 
     command_id = res['Command']['CommandId']
@@ -639,8 +634,7 @@
         while True:
             try:
                 res = ssmclient.get_command_invocation(
-                    CommandId=command_id,
-                    InstanceId=instance.id,
+                    CommandId=command_id, InstanceId=instance.id,
                 )
             except botocore.exceptions.ClientError as e:
                 if e.response['Error']['Code'] == 'InvocationDoesNotExist':
@@ -655,8 +649,9 @@
             elif res['Status'] in ('Pending', 'InProgress', 'Delayed'):
                 time.sleep(2)
             else:
-                raise Exception('command failed on %s: %s' % (
-                    instance.id, res['Status']))
+                raise Exception(
+                    'command failed on %s: %s' % (instance.id, res['Status'])
+                )
 
 
 @contextlib.contextmanager
@@ -711,10 +706,12 @@
     config['IamInstanceProfile'] = {
         'Name': 'hg-ephemeral-ec2-1',
     }
-    config.setdefault('TagSpecifications', []).append({
-        'ResourceType': 'instance',
-        'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}],
-    })
+    config.setdefault('TagSpecifications', []).append(
+        {
+            'ResourceType': 'instance',
+            'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}],
+        }
+    )
     config['UserData'] = WINDOWS_USER_DATA % password
 
     with temporary_ec2_instances(c.ec2resource, config) as instances:
@@ -723,7 +720,9 @@
         print('waiting for Windows Remote Management service...')
 
         for instance in instances:
-            client = wait_for_winrm(instance.public_ip_address, 'Administrator', password)
+            client = wait_for_winrm(
+                instance.public_ip_address, 'Administrator', password
+            )
             print('established WinRM connection to %s' % instance.id)
             instance.winrm_client = client
 
@@ -748,14 +747,17 @@
     # Store a reference to a good image so it can be returned one the
     # image state is reconciled.
     images = ec2resource.images.filter(
-        Filters=[{'Name': 'name', 'Values': [name]}])
+        Filters=[{'Name': 'name', 'Values': [name]}]
+    )
 
     existing_image = None
 
     for image in images:
         if image.tags is None:
-            print('image %s for %s lacks required tags; removing' % (
-                image.id, image.name))
+            print(
+                'image %s for %s lacks required tags; removing'
+                % (image.id, image.name)
+            )
             remove_ami(ec2resource, image)
         else:
             tags = {t['Key']: t['Value'] for t in image.tags}
@@ -763,15 +765,18 @@
             if tags.get('HGIMAGEFINGERPRINT') == fingerprint:
                 existing_image = image
             else:
-                print('image %s for %s has wrong fingerprint; removing' % (
-                      image.id, image.name))
+                print(
+                    'image %s for %s has wrong fingerprint; removing'
+                    % (image.id, image.name)
+                )
                 remove_ami(ec2resource, image)
 
     return existing_image
 
 
-def create_ami_from_instance(ec2client, instance, name, description,
-                             fingerprint):
+def create_ami_from_instance(
+    ec2client, instance, name, description, fingerprint
+):
     """Create an AMI from a running instance.
 
     Returns the ``ec2resource.Image`` representing the created AMI.
@@ -779,29 +784,19 @@
     instance.stop()
 
     ec2client.get_waiter('instance_stopped').wait(
-        InstanceIds=[instance.id],
-        WaiterConfig={
-            'Delay': 5,
-        })
+        InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
+    )
     print('%s is stopped' % instance.id)
 
-    image = instance.create_image(
-        Name=name,
-        Description=description,
-    )
+    image = instance.create_image(Name=name, Description=description,)
 
-    image.create_tags(Tags=[
-        {
-            'Key': 'HGIMAGEFINGERPRINT',
-            'Value': fingerprint,
-        },
-    ])
+    image.create_tags(
+        Tags=[{'Key': 'HGIMAGEFINGERPRINT', 'Value': fingerprint,},]
+    )
 
     print('waiting for image %s' % image.id)
 
-    ec2client.get_waiter('image_available').wait(
-        ImageIds=[image.id],
-    )
+    ec2client.get_waiter('image_available').wait(ImageIds=[image.id],)
 
     print('image %s available as %s' % (image.id, image.name))
 
@@ -827,9 +822,7 @@
         ssh_username = 'admin'
     elif distro == 'debian10':
         image = find_image(
-            ec2resource,
-            DEBIAN_ACCOUNT_ID_2,
-            'debian-10-amd64-20190909-10',
+            ec2resource, DEBIAN_ACCOUNT_ID_2, 'debian-10-amd64-20190909-10',
         )
         ssh_username = 'admin'
     elif distro == 'ubuntu18.04':
@@ -871,10 +864,12 @@
         'SecurityGroupIds': [c.security_groups['linux-dev-1'].id],
     }
 
-    requirements2_path = (pathlib.Path(__file__).parent.parent /
-                          'linux-requirements-py2.txt')
-    requirements3_path = (pathlib.Path(__file__).parent.parent /
-                          'linux-requirements-py3.txt')
+    requirements2_path = (
+        pathlib.Path(__file__).parent.parent / 'linux-requirements-py2.txt'
+    )
+    requirements3_path = (
+        pathlib.Path(__file__).parent.parent / 'linux-requirements-py3.txt'
+    )
     with requirements2_path.open('r', encoding='utf-8') as fh:
         requirements2 = fh.read()
     with requirements3_path.open('r', encoding='utf-8') as fh:
@@ -882,12 +877,14 @@
 
     # Compute a deterministic fingerprint to determine whether image needs to
     # be regenerated.
-    fingerprint = resolve_fingerprint({
-        'instance_config': config,
-        'bootstrap_script': BOOTSTRAP_DEBIAN,
-        'requirements_py2': requirements2,
-        'requirements_py3': requirements3,
-    })
+    fingerprint = resolve_fingerprint(
+        {
+            'instance_config': config,
+            'bootstrap_script': BOOTSTRAP_DEBIAN,
+            'requirements_py2': requirements2,
+            'requirements_py3': requirements3,
+        }
+    )
 
     existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)
 
@@ -902,9 +899,11 @@
         instance = instances[0]
 
         client = wait_for_ssh(
-            instance.public_ip_address, 22,
+            instance.public_ip_address,
+            22,
             username=ssh_username,
-            key_filename=str(c.key_pair_path_private('automation')))
+            key_filename=str(c.key_pair_path_private('automation')),
+        )
 
         home = '/home/%s' % ssh_username
 
@@ -926,8 +925,9 @@
                 fh.chmod(0o0700)
 
             print('executing bootstrap')
-            chan, stdin, stdout = ssh_exec_command(client,
-                                                   '%s/bootstrap' % home)
+            chan, stdin, stdout = ssh_exec_command(
+                client, '%s/bootstrap' % home
+            )
             stdin.close()
 
             for line in stdout:
@@ -937,17 +937,28 @@
             if res:
                 raise Exception('non-0 exit from bootstrap: %d' % res)
 
-            print('bootstrap completed; stopping %s to create %s' % (
-                  instance.id, name))
+            print(
+                'bootstrap completed; stopping %s to create %s'
+                % (instance.id, name)
+            )
 
-        return create_ami_from_instance(ec2client, instance, name,
-                                        'Mercurial Linux development environment',
-                                        fingerprint)
+        return create_ami_from_instance(
+            ec2client,
+            instance,
+            name,
+            'Mercurial Linux development environment',
+            fingerprint,
+        )
 
 
 @contextlib.contextmanager
-def temporary_linux_dev_instances(c: AWSConnection, image, instance_type,
-                                  prefix='hg-', ensure_extra_volume=False):
+def temporary_linux_dev_instances(
+    c: AWSConnection,
+    image,
+    instance_type,
+    prefix='hg-',
+    ensure_extra_volume=False,
+):
     """Create temporary Linux development EC2 instances.
 
     Context manager resolves to a list of ``ec2.Instance`` that were created
@@ -979,8 +990,9 @@
 
     # This is not an exhaustive list of instance types having instance storage.
     # But
-    if (ensure_extra_volume
-        and not instance_type.startswith(tuple(INSTANCE_TYPES_WITH_STORAGE))):
+    if ensure_extra_volume and not instance_type.startswith(
+        tuple(INSTANCE_TYPES_WITH_STORAGE)
+    ):
         main_device = block_device_mappings[0]['DeviceName']
 
         if main_device == 'xvda':
@@ -988,17 +1000,20 @@
         elif main_device == '/dev/sda1':
             second_device = '/dev/sdb'
         else:
-            raise ValueError('unhandled primary EBS device name: %s' %
-                             main_device)
+            raise ValueError(
+                'unhandled primary EBS device name: %s' % main_device
+            )
 
-        block_device_mappings.append({
-            'DeviceName': second_device,
-            'Ebs': {
-                'DeleteOnTermination': True,
-                'VolumeSize': 8,
-                'VolumeType': 'gp2',
+        block_device_mappings.append(
+            {
+                'DeviceName': second_device,
+                'Ebs': {
+                    'DeleteOnTermination': True,
+                    'VolumeSize': 8,
+                    'VolumeType': 'gp2',
+                },
             }
-        })
+        )
 
     config = {
         'BlockDeviceMappings': block_device_mappings,
@@ -1019,9 +1034,11 @@
 
         for instance in instances:
             client = wait_for_ssh(
-                instance.public_ip_address, 22,
+                instance.public_ip_address,
+                22,
                 username='hg',
-                key_filename=ssh_private_key_path)
+                key_filename=ssh_private_key_path,
+            )
 
             instance.ssh_client = client
             instance.ssh_private_key_path = ssh_private_key_path
@@ -1033,8 +1050,9 @@
                 instance.ssh_client.close()
 
 
-def ensure_windows_dev_ami(c: AWSConnection, prefix='hg-',
-                           base_image_name=WINDOWS_BASE_IMAGE_NAME):
+def ensure_windows_dev_ami(
+    c: AWSConnection, prefix='hg-', base_image_name=WINDOWS_BASE_IMAGE_NAME
+):
     """Ensure Windows Development AMI is available and up-to-date.
 
     If necessary, a modern AMI will be built by starting a temporary EC2
@@ -1100,13 +1118,15 @@
 
     # Compute a deterministic fingerprint to determine whether image needs
     # to be regenerated.
-    fingerprint = resolve_fingerprint({
-        'instance_config': config,
-        'user_data': WINDOWS_USER_DATA,
-        'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL,
-        'bootstrap_commands': commands,
-        'base_image_name': base_image_name,
-    })
+    fingerprint = resolve_fingerprint(
+        {
+            'instance_config': config,
+            'user_data': WINDOWS_USER_DATA,
+            'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL,
+            'bootstrap_commands': commands,
+            'base_image_name': base_image_name,
+        }
+    )
 
     existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)
 
@@ -1131,9 +1151,7 @@
             ssmclient,
             [instance],
             'AWS-RunPowerShellScript',
-            {
-                'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),
-            },
+            {'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),},
         )
 
         # Reboot so all updates are fully applied.
@@ -1145,10 +1163,8 @@
         print('rebooting instance %s' % instance.id)
         instance.stop()
         ec2client.get_waiter('instance_stopped').wait(
-            InstanceIds=[instance.id],
-            WaiterConfig={
-                'Delay': 5,
-            })
+            InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
+        )
 
         instance.start()
         wait_for_ip_addresses([instance])
@@ -1159,8 +1175,11 @@
         # TODO figure out a workaround.
 
         print('waiting for Windows Remote Management to come back...')
-        client = wait_for_winrm(instance.public_ip_address, 'Administrator',
-                                c.automation.default_password())
+        client = wait_for_winrm(
+            instance.public_ip_address,
+            'Administrator',
+            c.automation.default_password(),
+        )
         print('established WinRM connection to %s' % instance.id)
         instance.winrm_client = client
 
@@ -1168,14 +1187,23 @@
         run_powershell(instance.winrm_client, '\n'.join(commands))
 
         print('bootstrap completed; stopping %s to create image' % instance.id)
-        return create_ami_from_instance(ec2client, instance, name,
-                                        'Mercurial Windows development environment',
-                                        fingerprint)
+        return create_ami_from_instance(
+            ec2client,
+            instance,
+            name,
+            'Mercurial Windows development environment',
+            fingerprint,
+        )
 
 
 @contextlib.contextmanager
-def temporary_windows_dev_instances(c: AWSConnection, image, instance_type,
-                                    prefix='hg-', disable_antivirus=False):
+def temporary_windows_dev_instances(
+    c: AWSConnection,
+    image,
+    instance_type,
+    prefix='hg-',
+    disable_antivirus=False,
+):
     """Create a temporary Windows development EC2 instance.
 
     Context manager resolves to the list of ``EC2.Instance`` that were created.
@@ -1205,6 +1233,7 @@
             for instance in instances:
                 run_powershell(
                     instance.winrm_client,
-                    'Set-MpPreference -DisableRealtimeMonitoring $true')
+                    'Set-MpPreference -DisableRealtimeMonitoring $true',
+                )
 
         yield instances