Attempts to fix inconsistencies with terragrunt local tfstate, s3 backend state and dynamodb lock state

fix_terragrunt_state.py

#!/usr/bin/env python3

import os
import json

from pprint import pprint

import boto3


class State:
    def __init__(self, local_filename):
        self.s3 = boto3.resource('s3')
        self.dynamodb = boto3.resource('dynamodb')

        self.local_data = json.load(open(local_filename))
        self._load_from_backend()

    def _load_from_backend(self):
        # only s3 supported
        assert safeget(self.local_data, 'backend', 'type') == 's3'

        backend_config = self.local_data['backend']['config']
        bucket = backend_config['bucket']
        s3_key = backend_config['key']

        self.dyn_table = self.dynamodb.Table(backend_config['dynamodb_table'])
        self.dyn_active_key = {'LockID': '{}/{}'.format(bucket, s3_key)}
        self.dyn_state_key = {'LockID': '{}/{}-md5'.format(bucket, s3_key)}

        self.s3_obj = self.s3.Bucket(bucket).Object(s3_key)
        self.state = self.load_s3_state()
        self.active_lock = self.load_lock(self.dyn_active_key)
        self.state_lock = self.load_lock(self.dyn_state_key)


    def load_s3_state(self):
        print("Loading state: {}".format(self.s3_obj))
        try:
            data = self.s3_obj.get()['Body'].read().decode('utf-8')
            pprint(data)
            return data
        except:
            print("Not found")
            return None

    def empty(self):
        return False # TODO

    def delete_state(self):
        print("Deleting state: {}".format(self.s3_obj))
        return self.s3_obj.delete()

    def load_lock(self, key):
        print("Loading lock {}".format(key))
        try:
            resp = self.dyn_table.get_item(Key=key)
            lock = resp['Item']
            pprint(lock)
            return lock
        except:
            print("Not found")
            return None

    def delete_lock(self, key):
        print("Deleting lock {}".format(key))
        return self.dyn_table.delete_item(Key=key)

    def fix(self):
        if self.state:
            if self.empty():
                print("Deleting state & lock as state is empty")
                self.delete_state()
                self.delete_lock(self.dyn_state_key)
            else:
                print("State is not empty, not cleaning up.")
        else:
            print("Deleting state lock as state is missing")
            self.delete_lock(self.dyn_state_key)

        if self.active_lock:
            self.delete_lock(self.dyn_active_key)

def safeget(dct, *keys):
    dct = dict(dct)
    for key in keys:
        try:
            dct = dct[key]
        except (KeyError, AttributeError, TypeError):
            return None
    return dct

def get_state_files():
    for folder, _, file in os.walk('.terragrunt-cache'):
        for filename in file:
            if filename.endswith('.tfstate'):
                yield os.path.join(folder, filename)


if __name__ == '__main__':
    files = list(get_state_files())
    assert len(files) == 1
    print("Local state: {}".format(files[0]))
    state = State(files[0])

    state.fix()