sdp_snapshot.py #1

#!/usr/bin/env python3
"""Create a coordinated cloud backup set, then expire only owned complete sets.

Run through snapshot.sh to initialize the SDP environment and journal metadata.
Only standard-library modules and the provider's authenticated CLI are required.
"""
import argparse
from datetime import datetime, timedelta, timezone
import hashlib
import json
import logging
import os
import re
import subprocess
import sys
import uuid

MANAGER = 'sdp-cloud-snapshot'
SCHEMA = '1'
RUN_PATTERN = r'\d{8}T\d{6}Z-[0-9a-f]{12}'
POINT_PATTERN = re.compile(r'sdp-v1-j[0-9]+-' + RUN_PATTERN + r'\Z')


class SnapshotError(RuntimeError):
    pass


def require(condition, message):
    if not condition:
        raise SnapshotError(message)


def command(argv):
    """Structured arguments and JSON responses; no shell or credential logging."""
    try:
        result = subprocess.run(argv, capture_output=True, text=True, check=False)
    except OSError as exc:
        raise SnapshotError('Could not start cloud CLI: ' + str(exc)) from None
    if result.stderr.strip():
        logging.warning('%s', result.stderr.strip())
    if result.returncode:
        raise SnapshotError('Cloud command failed (exit %d): %s' %
                            (result.returncode, result.stderr.strip() or result.stdout.strip()))
    if not result.stdout.strip():
        return None
    try:
        return json.loads(result.stdout)
    except ValueError:
        raise SnapshotError('Cloud command returned invalid JSON') from None


def timestamp(value):
    require(isinstance(value, str) and bool(value), 'Missing snapshot creation time')
    try:
        result = datetime.fromisoformat(value.replace('Z', '+00:00'))
    except ValueError:
        raise SnapshotError('Invalid snapshot creation time: ' + value) from None
    require(result.tzinfo is not None, 'Snapshot creation time must include a timezone')
    return result.astimezone(timezone.utc)


def positive_int(value):
    try:
        result = int(value)
    except (TypeError, ValueError):
        raise SnapshotError('Retention/wait interval must be a positive integer') from None
    require(result > 0, 'Retention/wait interval must be a positive integer')
    return result


def tags_dict(items):
    require(isinstance(items, list), 'Malformed AWS tags')
    result = {}
    for item in items:
        require(isinstance(item, dict) and isinstance(item.get('Key'), str), 'Malformed AWS tag')
        require(item['Key'] not in result, 'Duplicate AWS tag key')
        result[item['Key']] = item.get('Value')
    return result


def owned(tags, expected):
    return isinstance(tags, dict) and all(tags.get(key) == value for key, value in expected.items())


def properties(resource):
    """Support both ARM properties objects and flattened Azure CLI responses."""
    require(isinstance(resource, dict), 'Malformed Azure resource')
    return resource.get('properties', resource)


class Backup:
    def __init__(self, instance, host, journal, retention=90, run_id=None, now=None, runner=command, environ=None):
        require(bool(instance) and bool(host), 'SDP instance and host are required')
        require(re.fullmatch(r'[0-9]+', str(journal)) is not None, 'Journal must be numeric')
        self.instance, self.host, self.journal = str(instance), host, str(journal)
        self.now = now or datetime.now(timezone.utc)
        self.run_id = run_id or self.now.strftime('%Y%m%dT%H%M%SZ-') + uuid.uuid4().hex[:12]
        require(re.fullmatch(RUN_PATTERN, self.run_id) is not None, 'Invalid backup run identifier')
        self.cutoff = self.now - timedelta(days=positive_int(retention))
        self.runner = runner
        self.env = os.environ if environ is None else environ

    def scope(self, account):
        return {'SDPManagedBy': MANAGER, 'SDPSchema': SCHEMA,
                'SDPAccount': account, 'SDPHost': self.host, 'SDPInstance': self.instance}

    def aws(self, *args):
        return self.runner(['aws', *args, '--output', 'json'])

    def run_aws(self):
        account = self.aws('sts', 'get-caller-identity')['Account']
        require(isinstance(account, str) and re.fullmatch(r'[0-9]{12}', account), 'Invalid AWS account identity')
        selected = []
        source_instance = None
        for suffix in ('root', 'hxdepots', 'hxlogs'):
            name = self.host + '-' + suffix
            response = self.aws('ec2', 'describe-volumes', '--filters',
                                json.dumps([{'Name': 'tag:Name', 'Values': [name]}]))
            volumes = response.get('Volumes')
            require(isinstance(volumes, list) and len(volumes) == 1,
                    'Expected exactly one volume named ' + name)
            volume = volumes[0]
            require(tags_dict(volume.get('Tags', [])).get('Name') == name,
                    'Volume Name tag differs from requested name ' + name)
            volume_id = volume.get('VolumeId', '')
            require(re.fullmatch(r'vol-[0-9a-f]+', volume_id) is not None and volume_id not in selected,
                    'Invalid or duplicate selected volume')
            attachments = volume.get('Attachments', [])
            require(len(attachments) == 1 and attachments[0].get('State') == 'attached',
                    'Selected volumes must each be attached to exactly one instance')
            instance_id = attachments[0].get('InstanceId', '')
            require(re.fullmatch(r'i-[0-9a-f]+', instance_id) is not None, 'Invalid EC2 instance ID')
            if source_instance is None:
                source_instance = instance_id
            require(instance_id == source_instance, 'Selected volumes span EC2 instances')
            selected.append(volume_id)
        expected_instance = self.env.get('SDP_AWS_INSTANCE_ID')
        require(not expected_instance or expected_instance == source_instance,
                'Selected volumes do not belong to SDP_AWS_INSTANCE_ID')
        response = self.aws('ec2', 'describe-instances', '--instance-ids', source_instance)
        instances = [vm for reservation in response.get('Reservations', []) for vm in reservation.get('Instances', [])]
        require(len(instances) == 1 and instances[0].get('InstanceId') == source_instance,
                'Could not identify the selected EC2 instance')
        vm = instances[0]
        devices = vm.get('BlockDeviceMappings', [])
        attached = [item['Ebs']['VolumeId'] for item in devices if 'Ebs' in item]
        roots = [item['Ebs']['VolumeId'] for item in devices
                 if 'Ebs' in item and item.get('DeviceName') == vm.get('RootDeviceName')]
        require(len(roots) == 1 and len(attached) == len(set(attached)) and set(selected) <= set(attached),
                'Ambiguous or changed EC2 volume attachment mapping')
        excluded = sorted(set(attached) - set(selected) - set(roots))
        require(len(excluded) <= 40, 'AWS supports at most 40 excluded data volumes')
        spec = {'InstanceId': source_instance, 'ExcludeBootVolume': roots[0] not in selected}
        if excluded:
            spec['ExcludeDataVolumeIds'] = excluded
        scope = dict(self.scope(account), SDPSourceInstance=source_instance)
        tags = dict(scope, SDPRun=self.run_id, SDPJournal=self.journal, SDPSetSize=str(len(selected)),
                    Contents='Perforce Checkpoints and Archives', Host=self.host,
                    Name=self.host + '-sdp-' + self.instance + '-j' + self.journal)
        tag_spec = [{'ResourceType': 'snapshot', 'Tags': [{'Key': key, 'Value': value} for key, value in tags.items()]}]
        logging.info('Creating AWS snapshot set %s for %s, journal %s', self.run_id, source_instance, self.journal)
        response = self.aws('ec2', 'create-snapshots', '--instance-specification', json.dumps(spec),
                            '--tag-specifications', json.dumps(tag_spec), '--description',
                            'SDP %s host %s journal %s run %s' % (self.instance, self.host, self.journal, self.run_id))
        snapshots = response.get('Snapshots', [])
        ids = self.validate_aws_set(snapshots, selected, account, tags, completed=False)
        logging.info('Snapshot set requested: %s', ', '.join(ids))
        self.aws('ec2', 'wait', 'snapshot-completed', '--snapshot-ids', *ids)
        final = self.aws('ec2', 'describe-snapshots', '--owner-ids', 'self', '--snapshot-ids', *ids)
        self.validate_aws_set(final.get('Snapshots', []), selected, account, tags, completed=True)
        logging.info('Snapshot set completed: %s', ', '.join(ids))
        self.expire_aws(scope, account)

    def validate_aws_set(self, snapshots, volumes, account, tags, completed):
        require(isinstance(snapshots, list) and len(snapshots) == len(volumes), 'Incomplete AWS snapshot set; expiration blocked')
        require({snap.get('VolumeId') for snap in snapshots} == set(volumes), 'AWS snapshot volume set changed; expiration blocked')
        ids = []
        for snap in snapshots:
            sid = snap.get('SnapshotId', '')
            require(re.fullmatch(r'snap-[0-9a-f]+', sid) is not None and sid not in ids, 'Invalid or duplicate AWS snapshot ID')
            require(snap.get('OwnerId') == account and owned(tags_dict(snap.get('Tags', [])), tags),
                    'AWS snapshot ownership/tags could not be verified; expiration blocked')
            require(snap.get('State') in (('completed',) if completed else ('pending', 'completed')),
                    'AWS snapshot did not complete successfully; expiration blocked')
            ids.append(sid)
        return ids

    def expire_aws(self, scope, account):
        filters = [{'Name': 'tag:' + key, 'Values': [value]} for key, value in scope.items()]
        response = self.aws('ec2', 'describe-snapshots', '--owner-ids', 'self', '--filters', json.dumps(filters))
        groups = {}
        for snap in response.get('Snapshots', []):
            tags = tags_dict(snap.get('Tags', []))
            if snap.get('OwnerId') != account or not owned(tags, scope):
                continue
            run_id = tags.get('SDPRun', '')
            if not re.fullmatch(RUN_PATTERN, run_id) or run_id == self.run_id:
                continue
            groups.setdefault(run_id, []).append((snap, tags))
        expired = []
        for run_id, records in groups.items():
            # Incomplete or partly failed sets are retained for manual inspection.
            if any(snap.get('State') != 'completed' or tags.get('SDPSetSize') != str(len(records))
                   for snap, tags in records):
                continue
            if len({snap.get('VolumeId') for snap, _ in records}) != len(records):
                continue
            if all(timestamp(snap.get('StartTime')) <= self.cutoff for snap, _ in records):
                for snap, _ in records:
                    sid = snap.get('SnapshotId', '')
                    require(re.fullmatch(r'snap-[0-9a-f]+', sid) is not None, 'Invalid owned snapshot ID; expiration blocked')
                    expired.append(sid)
        # Validate the entire candidate list before deleting anything.
        for sid in expired:
            logging.info('Deleting expired owned snapshot %s', sid)
            try:
                self.aws('ec2', 'delete-snapshot', '--snapshot-id', sid)
            except SnapshotError as exc:
                raise SnapshotError('Could not delete expired snapshot %s: %s' % (sid, exc)) from None

    def az(self, *args):
        return self.runner(['az', *args, '--subscription', self.subscription, '--output', 'json', '--only-show-errors'])

    def azure_collection(self, collection, expected_id, source_id, scope):
        require(collection.get('id', '').lower() == expected_id.lower(), 'Unexpected Azure restore point collection ID')
        require(properties(collection).get('source', {}).get('id', '').lower() == source_id.lower(),
                'Azure collection belongs to another VM; refusing to modify it')
        require(owned(collection.get('tags'), scope), 'Azure collection ownership tags differ; refusing to modify it')
        require(properties(collection).get('provisioningState') == 'Succeeded', 'Azure collection is not ready')

    def run_azure(self):
        def setting(name):
            value = self.env.get(name, '')
            require(bool(value), name + ' must be configured explicitly for Azure')
            return value
        self.subscription = setting('SDP_AZURE_SUBSCRIPTION')
        group, vm_name = setting('SDP_AZURE_RESOURCE_GROUP'), setting('SDP_AZURE_VM')
        timeout = str(positive_int(self.env.get('SDP_AZURE_WAIT_SECONDS', '3600')))
        account = self.az('account', 'show')
        self.subscription = account.get('id', '')
        require(re.fullmatch(r'[0-9a-fA-F-]{36}', self.subscription) is not None, 'Invalid Azure subscription identity')
        vm = self.az('vm', 'show', '--resource-group', group, '--name', vm_name)
        source_id = '/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/virtualMachines/%s' % (self.subscription, group, vm_name)
        require(vm.get('id', '').lower() == source_id.lower(), 'Azure VM identity differs from configured scope')
        vp = properties(vm)
        vm_uuid = vp.get('vmId')
        require(isinstance(vm_uuid, str) and bool(vm_uuid), 'Azure VM unique ID missing')
        storage = vp.get('storageProfile', {})
        os_disk = storage.get('osDisk', {})
        require(os_disk.get('osType') == 'Linux', 'This SDP snapshot entry point supports Linux Azure VMs')
        require(not vp.get('virtualMachineScaleSet'), 'Scale-set VMs require a separately validated restore-point workflow')
        disks = [os_disk] + storage.get('dataDisks', [])
        disk_ids = []
        for disk in disks:
            disk_id = disk.get('managedDisk', {}).get('id')
            require(bool(disk_id) and not disk.get('diffDiskSettings') and not disk.get('writeAcceleratorEnabled'),
                    'Crash-consistent Azure restore points require supported managed disks; ephemeral/write-accelerated disks are unsupported')
            detail = self.az('disk', 'show', '--ids', disk_id)
            require(detail.get('id', '').lower() == disk_id.lower(), 'Azure disk identity mismatch')
            require(detail.get('sku', {}).get('name') not in ('UltraSSD_LRS', 'PremiumV2_LRS')
                    and (properties(detail).get('maxShares') or 1) == 1,
                    'Ultra, Premium SSD v2 and shared disks are unsupported for crash consistency')
            disk_ids.append(disk_id.lower())
        require(len(disk_ids) == len(set(disk_ids)), 'Duplicate Azure managed disk IDs')
        scope = dict(self.scope(self.subscription), SDPSourceVM=vm_uuid)
        digest = hashlib.sha256((source_id.lower() + '\0' + self.host + '\0' + self.instance).encode()).hexdigest()[:10]
        default_collection = 'sdp-' + re.sub(r'[^a-zA-Z0-9_-]', '-', self.host)[:25] + '-' + re.sub(r'[^a-zA-Z0-9_-]', '-', self.instance)[:15] + '-' + digest
        collection_name = self.env.get('SDP_AZURE_COLLECTION') or default_collection
        require(re.fullmatch(r'[a-zA-Z0-9_.-]{1,80}', collection_name) is not None, 'Invalid Azure restore point collection name')
        collection_id = '/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/restorePointCollections/%s' % (self.subscription, group, collection_name)
        coll_args = ['--resource-group', group, '--collection-name', collection_name]
        collections = self.az('restore-point', 'collection', 'list', '--resource-group', group)
        require(isinstance(collections, list), 'Azure collection listing is incomplete or malformed')
        matches = [item for item in collections if item.get('name', '').lower() == collection_name.lower()]
        require(len(matches) <= 1, 'Duplicate Azure collection names')
        if matches:
            collection = self.az('restore-point', 'collection', 'show', *coll_args)
        else:
            logging.info('Creating dedicated Azure collection %s', collection_name)
            collection = self.az('restore-point', 'collection', 'create', *coll_args,
                                 '--source-id', source_id, '--location', vm['location'],
                                 '--tags', *[key + '=' + value for key, value in scope.items()])
        self.azure_collection(collection, collection_id, source_id, scope)
        name = 'sdp-v1-j' + self.journal + '-' + self.run_id
        point_args = [*coll_args, '--name', name]
        logging.info('Creating Azure restore point %s for %s', name, source_id)
        self.az('restore-point', 'create', *point_args, '--consistency-mode', 'CrashConsistent')
        self.az('restore-point', 'wait', *point_args, '--created', '--timeout', timeout)
        # Provisioning success can precede completion of background disk copying.
        count = str(len(disk_ids))
        complete = 'length(instanceView.diskRestorePoints) == `' + count + '` && length(instanceView.diskRestorePoints[?replicationStatus.completionPercent == `100`]) == `' + count + '`'
        self.az('restore-point', 'wait', *point_args, '--expand', 'instanceView', '--custom', complete, '--timeout', timeout)
        point = self.az('restore-point', 'show', *point_args, '--instance-view')
        self.validate_azure_point(point, collection_id, vm_uuid, name, disk_ids)
        logging.info('Azure restore point completed: %s', point['id'])
        collection = self.az('restore-point', 'collection', 'show', *coll_args, '--restore-points')
        self.azure_collection(collection, collection_id, source_id, scope)
        candidates = properties(collection).get('restorePoints')
        require(isinstance(candidates, list), 'Azure restore point listing is incomplete')
        expired = []
        for candidate in candidates:
            old_name = candidate.get('name', '')
            if not POINT_PATTERN.fullmatch(old_name) or old_name == name:
                continue
            props = properties(candidate)
            if (props.get('provisioningState') != 'Succeeded'
                    or props.get('consistencyMode') != 'CrashConsistent'
                    or props.get('sourceMetadata', {}).get('vmId') != vm_uuid):
                continue
            expected_id = collection_id + '/restorePoints/' + old_name
            require(candidate.get('id', '').lower() == expected_id.lower(), 'Azure restore point lies outside owned collection')
            if timestamp(props.get('timeCreated')) <= self.cutoff:
                detail = self.az('restore-point', 'show', *coll_args, '--name', old_name, '--instance-view')
                self.validate_azure_point(detail, collection_id, vm_uuid, old_name)
                expired.append(old_name)
        for old_name in expired:
            logging.info('Deleting expired owned Azure restore point %s', old_name)
            self.az('restore-point', 'delete', *coll_args, '--name', old_name, '--yes')

    def validate_azure_point(self, point, collection_id, vm_uuid, name, disk_ids=None):
        require(point.get('id', '').lower() == (collection_id + '/restorePoints/' + name).lower(), 'Unexpected Azure restore point identity')
        props = properties(point)
        require(props.get('provisioningState') == 'Succeeded' and props.get('consistencyMode') == 'CrashConsistent',
                'Azure restore point is not a successful crash-consistent backup; expiration blocked')
        metadata = props.get('sourceMetadata', {})
        require(metadata.get('vmId') == vm_uuid, 'Azure restore point source VM differs; expiration blocked')
        profile = metadata.get('storageProfile', {})
        disks = [profile.get('osDisk', {})] + profile.get('dataDisks', [])
        source_disks = [disk.get('managedDisk', {}).get('id', '').lower() for disk in disks]
        restore_ids = [disk.get('diskRestorePoint', {}).get('id', '').lower() for disk in disks]
        require(all(source_disks) and all(restore_ids) and len(restore_ids) == len(set(restore_ids)),
                'Azure restore point disk set is incomplete')
        if disk_ids is not None:
            require(len(source_disks) == len(disk_ids) and set(source_disks) == set(disk_ids), 'Azure restore point disk set changed; expiration blocked')
        views = props.get('instanceView', {}).get('diskRestorePoints', [])
        require(len(views) == len(restore_ids) and {v.get('id', '').lower() for v in views} == set(restore_ids),
                'Azure disk copy status is incomplete; expiration blocked')
        require(all(view.get('replicationStatus', {}).get('completionPercent') == 100
                    and view.get('replicationStatus', {}).get('status', {}).get('code', '').lower() == 'replicationstate/succeeded'
                    for view in views), 'Azure disk copy has not completed; expiration blocked')


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--provider', choices=('aws', 'azure'), default='aws')
    parser.add_argument('--instance', required=True)
    parser.add_argument('--host', required=True)
    parser.add_argument('--journal', required=True)
    parser.add_argument('--retention-days', default=os.environ.get('SDP_SNAPSHOT_RETENTION_DAYS', '90'))
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s')
    try:
        backup = Backup(args.instance, args.host, args.journal, args.retention_days)
        logging.info('Starting %s backup run %s', args.provider, backup.run_id)
        getattr(backup, 'run_' + args.provider)()
    except (SnapshotError, KeyError, TypeError, ValueError) as exc:
        logging.error('Snapshot processing failed: %s', exc)
        return 1
    logging.info('Snapshot processing completed OK')
    return 0


if __name__ == '__main__':
    sys.exit(main())