from __future__ import print_function
import boto3
import os
import json

def lambda_handler(event, context):
    print("Received event: " + json.dumps(event, indent=2))
    if event["detail"]["event"] == "createSnapshot":
        handle_create_snapshot(event, context)
    elif event["detail"]["event"] == "shareSnapshot":
        handle_shared_snapshot(event, context)
    else:
        print("Unknown event")

def handle_create_snapshot(event, context):
    target = os.environ["TARGET_ACCOUNT_ID"]
    print("Share target is: '{0}'".format(target))
    if target == "DISABLED":
        return

    if event["detail"]["result"] == "succeeded":
        print("Sharing snapshot with destination {0}".format(target))
        ec2  = boto3.resource('ec2')
        # "snapshot_id": "arn:aws:ec2::us-west-2:snapshot/snap-01234567",
        snap_id = event["detail"]["snapshot_id"].split("/")[1]
        snap = ec2.Snapshot(snap_id)
        snap.modify_attribute(
                Attribute='createVolumePermission',
                CreateVolumePermission={
                    'Add': [{'UserId': str(target)}]
                    })


def handle_shared_snapshot(event, context):
    target_region = os.environ.get("TARGET_REGION")
    print("Copy target is: '{0}'".format(target_region))
    if target_region == "DISABLED":
        return

    if event["detail"]["result"] == "succeeded":
        ec2_source = boto3.resource('ec2')
        ec2_target = boto3.session.Session().client('ec2', region_name=target_region)
        # "snapshot_id": "arn:aws:ec2::us-west-2:snapshot/snap-01234567",
        snap_id = event["detail"]["snapshot_id"].split("/")[1]
        snap = ec2_source.Snapshot(snap_id)
        print("Copying snapshot {0} from {1} to region {2}".format(snap_id, event["account"], target_region))
        source_region = event["region"]
        # need to copy from the *target* region because snap.copy ignores destinationregion...
        ec2_target.copy_snapshot(
                SourceSnapshotId=snap_id,
                Description="Copy of {} from {} ({})".format(snap_id, snap.owner_id, snap.description),
                SourceRegion=source_region,
                )