s3.py

"""
Utils/mixins related to accessing the stack S3 bucket.
"""

import boto3
import botocore
import getpass
import hashlib
import json
import os
from . import base


class StackTracker(object):
    """ State and config for a stack.  This wraps the single object kept in
    the stacks S3 bucket. """

    default_region = 'us-west-2'

    def __init__(self, s3_obj):
        self._s3_obj = s3_obj
        self._meta_keys = set()
        self._body_keys = set()
        try:
            data = s3_obj.get()
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] != "NoSuchKey":
                raise
        else:
            for key, val in data['Metadata'].items():
                self._meta_keys.add(key)
                setattr(self, key, val)
            for key, val in json.loads(data['Body'].read().decode()).items():
                self._body_keys.add(key)
                setattr(self, key, val)
            assert not (self._meta_keys & self._body_keys), 'key collision'

    @property
    def region(self):
        try:
            return self._region
        except AttributeError:
            return self.default_region

    @region.setter
    def region(self, value):
        self._region = value

    def set(self, key, value, meta=False):
        """ Use for new object creation to track a key. """
        keys = self._meta_keys if meta else self._body_keys
        keys.add(key)
        setattr(self, key, value)

    def unset(self, key, meta=False):
        '''Use to remove keys from meta and regular keys'''
        if meta:
            self._meta_keys.discard(key)
        else:
            self._body_keys.discard(key)
        setattr(self, key, None)

    def gen_passphrase_hash(self, passphrase):
        return hashlib.sha256(passphrase.encode()).hexdigest()

    def save(self, changelog=None, passphrase=None):
        self.check_passphrase(passphrase)
        return self._save(changelog)

    def _save(self, changelog=None):
        if changelog is not None:
            self.changelog.append({
                "ts": base.localnow().isoformat(),
                "msg": changelog,
                "operator": os.environ["OPERATOR_IDENT"]
            })
        meta = dict((x, getattr(self, x)) for x in self._meta_keys)
        body = dict((x, getattr(self, x)) for x in self._body_keys)
        self._s3_obj.put(Body=json.dumps(body), Metadata=meta,
                         ContentType='application/json')

    def get_changelog(self):
        """ Parse old tuple format and newer dict format. """
        to_dict = lambda x: {
            "ts": x[0],
            "msg": x[1],
            "operator": '-'
        }
        return [entry if isinstance(entry, dict) else to_dict(entry)
                for entry in self.changelog]

    def check_passphrase(self, passphrase=None):
        if self._passphrase_hash:
            if passphrase is None:
                passphrase = getpass.getpass("Passphrase: ")
            passhash = self.gen_passphrase_hash(passphrase)
            if self._passphrase_hash != passhash:
                raise SystemExit("Invalid Passphrase")
        return passphrase

    def delete(self, passphrase=None):
        self.check_passphrase(passphrase)
        self._s3_obj.delete()

    def add_resource(self, kind, ident, **options):
        """ A resource is an AWS asset that we need to track. """
        if self.has_resource(kind, ident):
            raise ValueError('Duplicate resource: %s %s' % (kind, ident))
        self.resources.append({
            "kind": kind,
            "ident": ident,
            "options": options
        })
        self._save()

    def has_resource(self, kind, ident):
        for x in self.resources:
            if x['kind'] == kind and x['ident'] == ident:
                return True
        return False

    def update_state(self, state):
        self.state = state
        self._save()


class S3BucketMixin(object):

    stacks_bucket = 'homecu.odyssey.stacks'

    def prerun(self, *args, **kwargs):
        s3_res = boto3.resource('s3')
        self.stacks = s3_res.Bucket(self.stacks_bucket)
        super().prerun(*args, **kwargs)

    def get_stack(self, name):
        """ Load a tracker object from S3.  If the file does not exist just
        return None. """
        stack_obj = self.stacks.Object(name)
        try:
            stack_obj.load()
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] != "404":
                raise
        else:
            return StackTracker(stack_obj)