#!/usr/bin/env python

from pprint import pprint
from quickconfig import Extractor
import argparse
import getpass
import json
import time
import sys
import os

from pycloud.core.security import AESEncryption, KeyPair, get_agent_keys
from pycloud.core.cli import user_edit

class VaultFile():
    encoding = 'utf-8'

    def __init__(self, path, key=None):
        self.path = path
        self.key = key

    def metadata(self):
        contents = json.loads(self._raw_read())
        contents.pop('data', None) 
        return contents

    def _raw_read(self):
        with open(self.path, 'r') as f:
            return f.read()

    def _raw_write(self, data):
        with open(self.path, 'w') as f:
            return f.write(data)

    def read(self):
        metadata = json.loads(self._raw_read())
        encrypted_payload = metadata.pop('data')
        assert metadata['encryption_method'] == 'AES'
        payload = AESEncryption(self.key, encoding=self.encoding).decrypt(encrypted_payload)
        payload = self._data_load(payload)
        return metadata, payload

    def write(self, payload, extra_metadata=None):
        payload = self._data_prep(payload)
        payload = AESEncryption(self.key, encoding=self.encoding).encrypt(payload, encode=True)
        metadata = extra_metadata or {}
        metadata.update({
            'data': payload,
            'info': 'This file is a VaultFile generated by pycloud.',
            'encryption_method': 'AES',
        })
        self._raw_write(json.dumps(metadata, indent=4))

    def verify_key(self, key):
        return False

    def exists(self):
        """ Really dumb check to see if file exists. Race conditions are acceptable"""
        return os.path.exists(self.path)

    def _data_prep(self, data):
        return data

    def _data_load(self, data):
        return data

class JsonVaultFile(VaultFile):
    encoding = 'utf-8'

    def _data_prep(self, data):
        return json.dumps(data, indent=4)

    def _data_load(self, data):
        return json.loads(data)

class Vault():
    def __init__(self, data_file, private_key=None, key_db=None, user=None):
        self.data_file_location = self._path(data_file)
        self.private_key_path = self._path(private_key)
        self.key_db_location = self._path(key_db)
        self.user = user
        self.metadata = None
        self.data = None
        self.users = None
        self.loaded = False

    def _path(self, path):
        if path:
            return os.path.expanduser(path)
        else:
            return path

    def get_private_keys(self):
        if self.private_key_path:
            return [KeyPair(private_key_path=self.private_key_path)]
        else:
            return get_agent_keys()
        
    def get_encryption_key(self):
        if self.key_db_location:
            with open(self.key_db_location) as f:
                data = json.loads(f.read())
                user_db = data['users']
        else:
            self.load(metadata_only=True)
            user_db = self.metadata['users']

        user_data_by_name = {u['user']: u for u in user_db}
        user_data_by_pub_key = {u['public_key']: u for u in user_db}

        user_data = None
        user_key = None
        for key in self.get_private_keys():
            pub_key = key.public_key_str()
            user_data = user_data_by_pub_key.get(pub_key, None)
            user_key = key 
            if user_data is not None:
                break

        if user_data is None:
            return None
        else:
            return user_key.decrypt(user_data['encrypted_key'])

    def change_key(self):
        if not self.loaded:
            self.load()
        old_key = self.get_encryption_key()
        new_key = AESEncryption.generate_key()
        for user in self.users:
            user_key = KeyPair(pub_data=user['public_key'])
            user['encrypted_key'] = user_key.encrypt(new_key, encode_payload=True)
        self.save(encryption_key=new_key)

    def load(self, metadata_only=False):
        if metadata_only:
            data_file = JsonVaultFile(self.data_file_location)
            self.metadata = data_file.metadata()
        else:
            encryption_key = self.get_encryption_key()
            if not encryption_key:
                return False
            data_file = JsonVaultFile(self.data_file_location, encryption_key)
            self.metadata, self.data = data_file.read()
        if self.key_db_location:
            with open('r') as f:
                self.users = json.loads(f.read())['users']
        else:
            self.users = self.metadata['users']
        self.loaded = True
        return self.loaded

    def save(self, data=None, encryption_key=None):
        if data is not None:
            self.data = data

        if not self.loaded:
            return False

        metadata = self.metadata or {}
        if self.key_db_location:
            with open(self.key_db_location, 'w') as f:
                f.write(json.dumps({
                    'users': self.users
                }))
        else:
            metadata['users'] = self.users

        if encryption_key is None:
            encryption_key = self.get_encryption_key()
        data_file = JsonVaultFile(self.data_file_location, encryption_key)
        data_file.write(self.data, metadata)
        return True

    def _add_user(self, user, keypair, encryption_key):
        encrypted_key = keypair.encrypt(encryption_key, encode_payload=True)
        user = {
            'user': user,
            'public_key': keypair.public_key_str(),
            'encrypted_key': encrypted_key
        }
        if self.loaded:
            self.users.append(user)
        else:
            self.users = [user]
        return True

    def add_user(self, user, keypair):
        return self._add_user(user, keypair, self.get_encryption_key())

    def create(self, data=None):
        self.users = []
        self.loaded = True
        keypair = self.get_private_keys()[0]
        encryption_key = AESEncryption.generate_key()
        self._add_user(self.user, keypair, encryption_key)

        payload = data or {}
        metdata = None
        self.save(payload, encryption_key=encryption_key)

def vault_cli(source):
    parser = argparse.ArgumentParser(description='Create update')
    parser.add_argument('action', nargs='+')
    parser.add_argument('-d', '--vault', required=True)
    parser.add_argument('-k', '--key_db', default=None)
    parser.add_argument('-u', '--user', default=getpass.getuser())
    parser.add_argument('-i', '--private_key', default=None)
    args = parser.parse_args(source)
    action, action_args =  args.action[0], args.action[1:]

    vault = Vault(
        args.vault,
        key_db=args.key_db,
        user=args.user,
        private_key=args.private_key
    )
    if action == 'create': 
        vault.create()

    elif action == 'view': 
        success = vault.load()
        if not success:
            print('Failed to decrypt')
        else:
            pprint(vault.data)

    elif action == 'show_key': 
        success = vault.load()
        if not success:
            print('Failed to decrypt')
        else:
            print(vault.get_encryption_key())

    elif action == 'edit': 
        success = vault.load()
        if not success:
            print('Failed to decrypt')
            return False
        tmp_file = args.vault + '.editing'
        contents = json.dumps(vault.data, indent=4)
        success = False
        for i in range(5):
            contents = user_edit(tmp_file, contents)
            try:
                new_data = json.loads(contents)
            except Exception as e:
                print('Edit Result Error:', str(e))
                time.sleep(4)
            else:
                vault.data = new_data
                success = True
                break
        if success:
            vault.save()
        else:
            print('You failed too many times. Edits discarded.')

    elif action == 'add_user':
        success = vault.load()
        if not success:
            print('Failed to decrypt')
            return False
        user_name = action_args[0]
        user_key = KeyPair(private_key_path=os.path.expanduser(action_args[1]))
        vault.add_user(user_name, user_key)
        vault.save()

    elif action == 'rotate_key':
        success = vault.load()
        if not success:
            print('Failed to decrypt')
            return False
        vault.change_key()
        vault.save()

    elif action == 'set': 
        vault.load()
        pprint(vault.data)

    elif action == 'get': 
        vault.load()
        e = Extractor(vault.data)
        value_path = action_args[0]
        value = e.extract(value_path)
        print('::{}::'.format(value_path))
        pprint(value)

if __name__ == '__main__':
    source = sys.argv[1:]
    vault_cli(source)
