import boto3
import os,time
from botocore.exceptions import ClientError
from datetime import datetime, timedelta, timezone

client = boto3.client('ec2')
ec2 = boto3.resource('ec2')
ddb = boto3.resource('dynamodb')

# 'EBS-Snapshot-Lifecycle' is the table you store snapshot metadata, modify it to your own table name
table = ddb.Table('EBS-Snapshot-Lifecycle')

# set days to keep 快照保留周期
DAYS_TO_RETAIN = 7

def lambda_handler(event, context):

    os.environ['TZ'] = 'Asia/Shanghai'
    time.tzset()
    i=time.strftime('%X %x %Z')
    
    # set volume id, get volume who has a tag-key is 'Snapshot'
    describe_volumes=client.describe_volumes(
        Filters=[
            {
                'Name': 'tag-key',
                'Values': ['Snapshot',
                ]
                
            }
        ]
    )
    volume_id_list = []
    for vol in describe_volumes['Volumes']:
        volume_id_list.append(vol.get('VolumeId'))
    
    # set snapshot
    for volume_id in volume_id_list:
        volume = ec2.Volume(volume_id)
    
        for tags in volume.tags:
            if(tags.get('Key') == 'Name'):
                volume_name = tags.get('Value')
        description = volume_name + ' volume snapshot is created at ' + i
        
        try:
            response = client.create_snapshot(
                Description=description,
                VolumeId=volume_id)
        except:
            print('Create Snapshot occured error, Volume id is ' + volume_id)
        else:
            print('Snapshot is created succeed, Snapshot id is ' + response.get('SnapshotId'))
        
        # write into DynanoDB table
        snapshot_id = response.get('SnapshotId')
        start_time = response.get('StartTime')
        ttl_time = int(start_time.timestamp()) + 86400 * DAYS_TO_RETAIN
    
        try:
            table.put_item(
                Item={
                    'snapshot_id': snapshot_id,
                    'create_time': start_time.astimezone(timezone(timedelta(hours=8))).strftime('%X %x %Z'),
                    'description': description,
                    'timestamp': int(start_time.timestamp()),
                    'ttl': ttl_time,
                }
            )
        except:
            print('Write to DynanoDB Table occured error, snapshot id is ' + snapshot_id)
        else:
            print('Write to DynanoDB Table succeed, snapshot id is ' + snapshot_id)