
# Copyright 2016 Nitor Creations Oy
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function
import argparse
import locale
import os
import sys
import json
import argcomplete
import boto3
import requests
from requests.exceptions import ConnectionError
from .vault import Vault

SYS_ENCODING = locale.getpreferredencoding()
VAULT_STACK_VERSION = 20
TEMPLATE_STRING = """{
  "Parameters": {
    "paramBucketName": {
      "Default": "nitor-core-vault",
      "Type": "String",
      "Description": "Name of the vault bucket"
    }
  },
  "Resources": {
    "resourceDecryptRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "Path": "/",
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": "sts:AssumeRole",
              "Effect": "Allow",
              "Principal": {
                "Service": ["ec2.amazonaws.com"]
              }
            }
          ]
        }
      }
    },
    "resourceEncryptRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "Path": "/",
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": "sts:AssumeRole",
              "Effect": "Allow",
              "Principal": {
                "Service": ["ec2.amazonaws.com"]
              }
            }
          ]
        }
      }
    },
    "resourceLambdaRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "Path": "/",
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": "sts:AssumeRole",
              "Effect": "Allow",
              "Principal": {
                "Service": ["lambda.amazonaws.com", "edgelambda.amazonaws.com"]
              }
            }
          ]
        },
        "ManagedPolicyArns": ["arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"]
      }
    },    
    "kmsKey": {
      "Type": "AWS::KMS::Key",
      "Properties": {
        "KeyPolicy": {
          "Version": "2012-10-17",
          "Id": "key-default-2",
          "Statement": [
            {
              "Action": [
                "kms:*"
              ],
              "Principal": {
                "AWS": {
                  "Fn::Join": [
                    "",
                    [
                      "arn:aws:iam::",
                      {
                        "Ref": "AWS::AccountId"
                      },
                      ":root"
                    ]
                  ]
                }
              },
              "Resource": "*",
              "Effect": "Allow",
              "Sid": "allowAdministration"
            }
          ]
        },
        "Description": "Key for encrypting / decrypting secrets"
      }
    },
    "vaultBucket": {
      "Type": "AWS::S3::Bucket",
      "Properties": {
        "BucketName": {
          "Ref": "paramBucketName"
        }
      }
    },
    "iamPolicyEncrypt": {
      "Type": "AWS::IAM::ManagedPolicy",
      "Properties": {
        "PolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject"
              ],
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:aws:s3:::",
                    {
                      "Ref": "paramBucketName"
                    },
                    "/*"
                  ]
                ]
              },
              "Effect": "Allow",
              "Sid": "putVaultItems"
            },
            {
              "Action": [
                "s3:ListBucket"
              ],
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:aws:s3:::",
                    {
                      "Ref": "paramBucketName"
                    }
                  ]
                ]
              },
              "Effect": "Allow",
              "Sid": "listVault"
            },
            {
              "Action": [
                "cloudformation:DescribeStacks"
              ],
              "Resource": {
                "Ref": "AWS::StackId"
              },
              "Effect": "Allow",
              "Sid": "describeVault"
            },
            {
              "Action": [
                "kms:Decrypt",
                "kms:Encrypt",
                "kms:GenerateDataKey"
              ],
              "Resource": {
                "Fn::GetAtt": [
                  "kmsKey",
                  "Arn"
                ]
              },
              "Effect": "Allow",
              "Sid": "allowEncrypt"
            },
            {
              "Sid": "InvokeLambdaPermission",
              "Effect": "Allow",
              "Action": [
                  "lambda:InvokeFunction"
              ],
              "Resource": {"Fn::GetAtt": ["lambdaDecrypter", "Arn"]}
            }
          ]
        },
        "Description": "Policy to allow encrypting and decrypting vault secrets",
        "Roles": [
          {
            "Ref": "resourceEncryptRole"
          }
        ]
      }
    },
    "iamPolicyDecrypt": {
      "Type": "AWS::IAM::ManagedPolicy",
      "Properties": {
        "PolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "s3:GetObject"
              ],
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:aws:s3:::",
                    {
                      "Ref": "paramBucketName"
                    },
                    "/*"
                  ]
                ]
              },
              "Effect": "Allow",
              "Sid": "getVaultItems"
            },
            {
              "Action": [
                "s3:ListBucket"
              ],
              "Resource": {
                "Fn::Join": [
                  "",
                  [
                    "arn:aws:s3:::",
                    {
                      "Ref": "paramBucketName"
                    }
                  ]
                ]
              },
              "Effect": "Allow",
              "Sid": "listVault"
            },
            {
              "Action": [
                "cloudformation:DescribeStacks"
              ],
              "Resource": {
                "Ref": "AWS::StackId"
              },
              "Effect": "Allow",
              "Sid": "describeVault"
            },
            {
              "Action": [
                "kms:Decrypt"
              ],
              "Resource": {
                "Fn::GetAtt": [
                  "kmsKey",
                  "Arn"
                ]
              },
              "Effect": "Allow",
              "Sid": "allowDecrypt"
            },
            {
              "Sid": "InvokeLambdaPermission",
              "Effect": "Allow",
              "Action": [
                  "lambda:InvokeFunction"
              ],
              "Resource": {"Fn::GetAtt": ["lambdaDecrypter", "Arn"]}
            }
          ]
        },
        "Description": "Policy to allow decrypting vault secrets",
        "Roles": [
          {
            "Ref": "resourceDecryptRole"
          },
          {
            "Ref": "resourceLambdaRole"
          }
        ]
      }
    },
    "lambdaDecrypter": {
      "Type": "AWS::Lambda::Function",
      "Properties": {
        "Description": { "Fn::Sub": "Nitor Vault ${AWS::StackName} Decrypter"},
        "Handler": "index.handler",
        "MemorySize": 128,
        "Runtime": "python2.7",
        "Timeout": 300,
        "Role": {"Fn::GetAtt": ["resourceLambdaRole", "Arn"]},
        "FunctionName": {"Fn::Sub": "${AWS::StackName}-decrypter"},
        "Code": {
          "ZipFile" : { "Fn::Join" : ["\\n", [
            "import json",
            "import logging",
            "import boto3",
            "import base64",
            "from botocore.vendored import requests",
            "log = logging.getLogger()",
            "log.setLevel(logging.INFO)",
            "kms = boto3.client('kms')",
            "SUCCESS = 'SUCCESS'",
            "FAILED = 'FAILED'",
            "def handler(event, context):",
            "  ciphertext = event['ResourceProperties']['Ciphertext']",
            "  responseData = {}",
            "  try:",
            "    responseData['Plaintext'] = kms.decrypt(CiphertextBlob=base64.b64decode(ciphertext)).get('Plaintext')",
            "    log.info('Decrypt successful!')",
            "    send(event, context, SUCCESS, responseData, event['LogicalResourceId'])",
            "  except Exception as e:",
            "    error_msg = 'Failed to decrypt: ' + repr(e)",
            "    log.error(error_msg)",
            "    send(event, context, FAILED, responseData, event['LogicalResourceId'])",
            "    raise Exception(error_msg)",
            "",
            "def send(event, context, responseStatus, responseData, physicalResourceId):",
            "  responseUrl = event['ResponseURL']",
            "  responseBody = {}",
            "  responseBody['Status'] = responseStatus",
            "  responseBody['Reason'] = 'See the details in CloudWatch Log Stream: ' + context.log_stream_name",
            "  responseBody['PhysicalResourceId'] = physicalResourceId or context.log_stream_name",
            "  responseBody['StackId'] = event['StackId']",
            "  responseBody['RequestId'] = event['RequestId']",
            "  responseBody['LogicalResourceId'] = event['LogicalResourceId']",
            "  responseBody['Data'] = responseData",
            "  json_responseBody = json.dumps(responseBody)",
            "  headers = {",
            "    'content-type' : '',",
            "    'content-length' : str(len(json_responseBody))",
            "  }",
            "  try:",
            "    response = requests.put(responseUrl,",
            "                            data=json_responseBody,",
            "                            headers=headers)",
            "  except Exception as e:",
            "    log.warning('send(..) failed executing requests.put(..): ' + repr(e))"
          ]]}
        }
      }
    }
  },
  "Outputs": {
    "vaultBucketName": {
      "Description": "Vault Bucket",
      "Value": {
        "Ref": "vaultBucket"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "vaultBucketName"]]
        }
      }
    },
    "kmsKeyArn": {
      "Description": "KMS key Arn",
      "Value": {
        "Fn::GetAtt": [
          "kmsKey",
          "Arn"
        ]
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "kmsKeyArn"]]
        }
      }
    },
    "decryptRole": {
      "Description": "The role for decrypting",
      "Value": {
        "Ref": "resourceDecryptRole"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "decryptRole"]]
        }
      }
    },
    "encryptRole": {
      "Description": "The role for encrypting",
      "Value": {
        "Ref": "resourceEncryptRole"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "encryptRole"]]
        }
      }
    },
    "decryptPolicy": {
      "Description": "The policy for decrypting",
      "Value": {
        "Ref": "iamPolicyDecrypt"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "decryptPolicy"]]
        }
      }
    },
    "encryptPolicy": {
      "Description": "The policy for decrypting",
      "Value": {
        "Ref": "iamPolicyEncrypt"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "encryptPolicy"]]
        }
      }
    },
    "vaultStackVersion": {
      "Description": "The version of the currently deployed vault stack template",
      "Value": "%(version)s",
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "vaultStackVersion"]]
        }
      }
    },
    "lambdaDecrypterArn": {
      "Description": "Decrypter Lambda function ARN",
      "Value": {
        "Fn::Sub": "${lambdaDecrypter.Arn}"
      },
      "Export": {
        "Name": {
          "Fn::Join": [":", [{"Ref": "AWS::StackName"}, "lambdaDecrypterArn"]]
        }
      }
    }        
  }
}""" % {"version": VAULT_STACK_VERSION}

def _template():
    return json.dumps(json.loads(TEMPLATE_STRING))

def is_ec2():
    if sys.platform.startswith("win"):
        import wmi
        systeminfo = wmi.WMI().Win32_ComputerSystem()[0]
        return "EC2" == systeminfo.PrimaryOwnerName
    elif sys.platform.startswith("linux"):
        if os.path.isfile("/sys/hypervisor/uuid"):
            with open("/sys/hypervisor/uuid") as uuid:
                uuid_str = uuid.read()
                return uuid_str.startswith("ec2")
        else:
            return False

def region():
    """ Get default region - the region of the instance if run in an EC2 instance
    """
    # If it is set in the environment variable, use that
    if 'AWS_DEFAULT_REGION' in os.environ:
        return os.environ['AWS_DEFAULT_REGION']
    else:
        # Otherwise it might be configured in AWS credentials
        session = boto3.session.Session()
        if session.region_name:
            return session.region_name
        # If not configured and being called from an ec2 instance, use the
        # region of the instance
        elif is_ec2():
          try:
              response = requests.get('http://169.254.169.254/latest/dynamic/instance-identity/document')
              instance_data = json.loads(response.text)
              return instance_data['region']
          except ConnectionError:
              # no-op
              return ""
    return ""

def main():
    parser = argparse.ArgumentParser(description="Store and lookup locally " +\
                                     "encrypted data stored in S3")
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument('-s', '--store', help="Name of element to store. Opt" +\
                                              "ionally read from file name",
                        nargs='?', default="")
    action.add_argument('-l', '--lookup', help="Name of element to lookup")
    action.add_argument('-i', '--init', action='store_true',
                        help="Initializes a kms key and a s3 bucket with som" +\
                              "e roles for reading and writing on a fresh ac" +\
                              "count via CloudFormation. Means that the acco" +\
                              "unt used has to have rights to create the res" +\
                              "ources")
    action.add_argument('-u', '--update', action='store_true',
                        help="Updates the CloudFormation stack which declare" +\
                              "s all resources needed by the vault.")
    action.add_argument('-d', '--delete', help="Name of element to delete")
    action.add_argument('-a', '--all', action='store_true', help="List avail" +\
                                                                 "able secrets")
    parser.add_argument('-w', '--overwrite', action='store_true',
                        help="Add this argument if you want to overwrite an " +\
                             "existing element")
    store_data = parser.add_mutually_exclusive_group(required=False)
    store_data.add_argument('-v', '--value', help="Value to store")
    store_data.add_argument('-f', '--file', help="File to store. If no -s argument" +\
                                           " given, the name of the file is " +\
                                           "used as the default name. Give -" +\
                                           " for stdin")
    parser.add_argument('-o', "--outfile", help="The file to write the data to")
    parser.add_argument('-p', '--prefix', help="Optional prefix to store val" +\
                                               "ue under. empty by default")
    parser.add_argument('--vaultstack', help="Optional CloudFormation stack " +\
                                             "to lookup key and bucket. 'vau" +\
                                             "lt' by default")
    parser.add_argument('-b', '--bucket', help="Override the bucket name eit" +\
                                               "her for initialization or st" +\
                                               "oring and looking up values")
    parser.add_argument('-k', '--key-arn', help="Override the KMS key arn fo" +\
                                                "r storinig or looking up")
    parser.add_argument('--id', help="Give an IAM access key id to override " +\
                                     "those defined by environent")
    parser.add_argument('--secret', help="Give an IAM secret access key to o" +\
                                         "verride those defined by environent")
    parser.add_argument('-r', '--region', help="Give a region for the stack" +\
                                               "and bucket")
    argcomplete.autocomplete(parser)
    args = parser.parse_args()
    if args.store and not (args.value or args.file):
        parser.error("--store requires --value or --file")
    store_with_no_name = not args.store and not args.lookup and not args.init \
                         and not args.delete and not args.all and not args.update
    if store_with_no_name and not args.file:
        parser.error("--store requires a name or a --file argument to get the name to store")
    elif store_with_no_name:
        if args.file == "-":
            parser.error("--store requires a name for stdin")
        else:
            args.store = os.path.basename(args.file)
            data = open(args.file, 'rb').read()
    elif args.store:
        if args.value:
            data = args.value.encode()
        elif args.file == "-":
            data = sys.stdin.read()
        else:
            data = open(args.file, 'rb').read()
    if not args.vaultstack:
        if "VAULT_STACK" in os.environ:
            args.vaultstack = os.environ["VAULT_STACK"]
        else:
            args.vaultstack = "vault"

    if not args.bucket and "VAULT_BUCKET" in os.environ:
        args.bucket = os.environ["VAULT_BUCKET"]

    if not args.prefix and "VAULT_PREFIX" in os.environ:
        args.prefix = os.environ["VAULT_PREFIX"]
    elif not args.prefix:
        args.prefix = ""

    instance_data = None
    # Try to get region from instance metadata if not otherwise specified
    if not args.region:
        args.region = region()
    if args.region:
        os.environ['AWS_DEFAULT_REGION'] = args.region

    if not args.init and not args.update:
        vlt = Vault(vault_stack=args.vaultstack, vault_key=args.key_arn,
                    vault_bucket=args.bucket, vault_iam_id=args.id,
                    vault_iam_secret=args.secret, vault_prefix=args.prefix)
        if args.store:
            if args.overwrite or not vlt.exists(args.store):
                vlt.store(args.store, data)
            elif not args.overwrite:
                parser.error("Will not overwrite '" + args.store +
                             "' without the --overwrite (-w) flag")
        elif args.delete:
            vlt.delete(args.delete)
        elif args.all:
            data = vlt.all()
            if args.outfile and not args.outfile == "-":
                with open(args.outfile, 'w') as outf:
                    outf.write(data)
            else:
                sys.stdout.write(data)
        else:
            data = vlt.lookup(args.lookup)
            if args.outfile and not args.outfile == "-":
                with open(args.outfile, 'wb') as outf:
                    outf.write(data)
            else:
                sys.stdout.write(data.decode(SYS_ENCODING))
    else:
        if not args.bucket:
            sts = boto3.client("sts")
            account_id = sts.get_caller_identity()['Account']
            args.bucket = "vault-" + args.region + "-" + account_id
        clf = boto3.client("cloudformation")
        if args.init:
            try:
                clf.describe_stacks(StackName=args.vaultstack)
                print("Vault stack '" + args.vaultstack + "' already initialized")
            except:
                params = {}
                params['ParameterKey'] = "paramBucketName"
                params['ParameterValue'] = args.bucket
                clf.create_stack(StackName=args.vaultstack, TemplateBody=_template(),
                                 Parameters=[params], Capabilities=['CAPABILITY_IAM'])
        elif args.update:
            try:
                stack = clf.describe_stacks(StackName=args.vaultstack)
                deployed_version = None
                ok_to_update = False
                for output in stack['Stacks'][0]['Outputs']:
                    if output['OutputKey'] == 'vaultStackVersion':
                        deployed_version = int(output['OutputValue'])
                        if int(output['OutputValue']) < VAULT_STACK_VERSION:
                            ok_to_update = True
                        break
                if ok_to_update or deployed_version is None:
                    params = {}
                    params['ParameterKey'] = "paramBucketName"
                    params['UsePreviousValue'] = True
                    clf.update_stack(StackName=args.vaultstack, TemplateBody=_template(),
                                     Parameters=[params], Capabilities=['CAPABILITY_IAM'])
                else:
                    print("Current stack version %(cur_ver)s does not need update to " \
                          "version %(code_version)s" % {"cur_ver": deployed_version,
                                                        "code_version": VAULT_STACK_VERSION})
            except Exception as e:
                print("Error while updating stack '" + args.vaultstack + "': " + repr(e))
