import argparse
import datetime
import json

from collections import OrderedDict
from dateutil.tz import tzutc

import boto3


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--bucket', required=True,
                        help='S3 bucket containing state file')
    parser.add_argument('-o', '--object', required=True,
                        help='Object key of state file')
    parser.add_argument('-v', '--version', required=True,
                        help='Version of statefile object before import began')
    args = parser.parse_args()

    state_json_file = open('terraform.tfstate', 'r+')
    state_file = json.load(state_json_file, object_pairs_hook=OrderedDict)

    """
    Delete object versions that are newer than the starting version,
    as they contain sensitive information from import steps
    """
    client = boto3.client('s3')
    state_file_versions = client.list_object_versions(
        Bucket=args.bucket,
        Prefix=args.object,
    )

    # Get the modified time
    for version in state_file_versions['Versions']:
        if version['VersionId'] == args.version:
            version_datetime = version['LastModified']

    """
    Now remove any object versions created by
    the import process that have sensitive data
    """
    for version in state_file_versions['Versions']:
        if version['LastModified'] > version_datetime:
            print('Removing object version {}'.format(version['VersionId']))
            response = client.delete_object(
                Bucket=args.bucket,
                Key=args.object,
                VersionId=version['VersionId'],
            )
            print('Was the object version deleted? {}'.format(
                response))

    """
    Find all secret string resources.
    Data sources are kept up to date, so there is no reason to sanitize them
    as they would be overwritten on the next terraform plan
    """

    for resource in state_file['resources']:
        if resource['name'] == 'secret_string' and \
                resource['type'] == 'aws_secretsmanager_secret_version':
            resource['instances'][0]['attributes']['secret_string'] = \
                'REDACTED'

    # Increment serial number
    state_file['serial'] += 1

    state_json_file.seek(0)
    state_json_file.truncate()
    json.dump(state_file, state_json_file)


if __name__ == "__main__":
    main()