#!/usr/bin/python3
# *****************************************************************************
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#
# ******************************************************************************


import argparse
import os
import subprocess
from fabric import *
from invoke import task

parser = argparse.ArgumentParser()
parser.add_argument('--conf_service_base_name', type=str, help='unique name for DataLab environment')
parser.add_argument('--conf_network_type', type=str, default='',
                    help='Define in which network DataLab will be deployed. '
                         'Possible options: public|private')
parser.add_argument('--conf_vpc_cidr', type=str, default='', help='CIDR of VPC')
parser.add_argument('--conf_vpc2_cidr', type=str, default='', help='CIDR of secondary VPC')
parser.add_argument('--conf_allowed_ip_cidr', type=str, default='', help='Comma-separated CIDR of IPs which will have '
                                                                         'access to SSN')
parser.add_argument('--conf_user_subnets_range', type=str, default='', help='Range of subnets which will be using for '
                                                                            'users environments. For example: '
                                                                            '10.10.0.0/24 - 10.10.10.0/24')
parser.add_argument('--conf_private_subnet_prefix', type=str, default='24', help='Private subnet prefix')
parser.add_argument('--conf_additional_tags', type=str, default='', help='Additional tags in format '
                                                                         '"Key1:Value1;Key2:Value2"')
parser.add_argument('--conf_image_enabled', type=str, default='', help='Enable or Disable creating image at first time')
parser.add_argument('--aws_user_predefined_s3_policies', type=str, default='', help='Predefined policies for users '
                                                                                    'instances')
parser.add_argument('--aws_access_key', type=str, default='', help='AWS Access Key ID')
parser.add_argument('--aws_secret_access_key', type=str, default='', help='AWS Secret Access Key')
parser.add_argument('--aws_region', type=str, default='', help='AWS region')
parser.add_argument('--aws_zone', type=str, default='', help='AWS zone')
parser.add_argument('--azure_region', type=str, default='', help='Azure region')
parser.add_argument('--gcp_region', type=str, default='', help='GCP region')
parser.add_argument('--gcp_zone', type=str, default='', help='GCP zone')
parser.add_argument('--conf_os_family', type=str, default='',
                    help='Operating system type. Available options: debian, redhat')
parser.add_argument('--conf_cloud_provider', type=str, default='',
                    help='Where DataLab should be deployed. Available options: aws, azure, gcp')
parser.add_argument('--ssn_hosted_zone_name', type=str, default='', help='Name of hosted zone')
parser.add_argument('--ssn_hosted_zone_id', type=str, default='', help='ID of hosted zone')
parser.add_argument('--ssn_subdomain', type=str, default='', help='Subdomain name')
parser.add_argument('--ssn_assume_role_arn', type=str, default='', help='Role ARN for creating Route53 record in '
                                                                        'different AWS account')
parser.add_argument('--ssl_cert_path', type=str, default='', help='Full path to SSL certificate')
parser.add_argument('--ssl_key_path', type=str, default='', help='Full path to key for SSL certificate')
parser.add_argument('--aws_vpc_id', type=str, default='', help='AWS VPC ID')
parser.add_argument('--conf_duo_vpc_enable', type=str, default='false', help='Duo VPC scheme enable(true|false)')
parser.add_argument('--aws_vpc2_id', type=str, default='', help='Secondary AWS VPC ID')
parser.add_argument('--aws_peering_id', type=str, default='', help='Amazon peering connection id')
parser.add_argument('--azure_vpc_name', type=str, default='', help='Azure VPC Name')
parser.add_argument('--gcp_vpc_name', type=str, default='', help='GCP VPC Name')
parser.add_argument('--aws_subnet_id', type=str, default='', help='AWS Subnet ID')
parser.add_argument('--azure_subnet_name', type=str, default='', help='Azure Subnet Name')
parser.add_argument('--gcp_subnet_name', type=str, default='', help='GCP Subnet Name')
parser.add_argument('--aws_security_groups_ids', type=str, default='', help='One of more comma-separated Security '
                                                                            'groups IDs for SSN')
parser.add_argument('--azure_security_group_name', type=str, default='', help='One Security'
                                                                              'group name for SSN')
parser.add_argument('--azure_edge_security_group_name', type=str, default='', help='One Security '
                                                                              'group name for Edge node')
parser.add_argument('--gcp_firewall_name', type=str, default='', help='One of more comma-separated GCP Firewall rules '
                                                                      'for SSN')
parser.add_argument('--key_path', type=str, default='', help='Path to admin key (WITHOUT KEY NAME)')
parser.add_argument('--conf_key_name', type=str, default='', help='Admin key name (WITHOUT ".pem")')
parser.add_argument('--workspace_path', type=str, default='', help='Admin key name (WITHOUT ".pem")')
parser.add_argument('--conf_tag_resource_id', type=str, default='datalab', help='The name of user tag')
parser.add_argument('--conf_billing_tag', type=str, default='datalab', help='Billing tag')
parser.add_argument('--aws_ssn_instance_size', type=str, default='t2.large', help='The SSN instance shape')
parser.add_argument('--azure_ssn_instance_size', type=str, default='Standard_DS2_v2', help='The SSN instance shape')
parser.add_argument('--gcp_ssn_instance_size', type=str, default='n1-standard-2', help='The SSN instance shape')
parser.add_argument('--aws_account_id', type=str, default='', help='The ID of Amazon account')
parser.add_argument('--aws_billing_bucket', type=str, default='', help='The name of S3 bucket where billing reports '
                                                                       'will be placed.')
parser.add_argument('--aws_job_enabled', type=str, default='false', help='Billing format. Available options: '
                                                                         'true (aws), false(epam)')
parser.add_argument('--aws_report_path', type=str, default='', help='The path to billing reports directory in S3 '
                                                                    'bucket')
parser.add_argument('--azure_resource_group_name', type=str, default='', help='Name of Resource group in Azure')
parser.add_argument('--azure_auth_path', type=str, default='', help='Full path to Azure credentials JSON file')
parser.add_argument('--azure_datalake_enable', type=str, default='', help='Provision DataLake storage account')
parser.add_argument('--azure_ad_group_id', type=str, default='', help='ID of Azure AD group')
parser.add_argument('--azure_offer_number', type=str, default='', help='Azure offer number')
parser.add_argument('--azure_currency', type=str, default='', help='Azure currency code')
parser.add_argument('--azure_locale', type=str, default='', help='Azure locale')
parser.add_argument('--azure_application_id', type=str, default='', help='Azure login application ID')
parser.add_argument('--azure_validate_permission_scope', type=str, default='true', help='Azure permission scope '
                                                                                        'validation(true|false).')
parser.add_argument('--azure_oauth2_enabled', type=str, default='false', help='Using OAuth2 for logging in DataLab')
parser.add_argument('--azure_region_info', type=str, default='', help='Azure region info')
parser.add_argument('--azure_source_vpc_name', type=str, default='', help='Azure VPC source Name')
parser.add_argument('--azure_source_resource_group_name', type=str, default='', help='Azure source resource group')
parser.add_argument('--gcp_project_id', type=str, default='', help='The project ID in Google Cloud Platform')
parser.add_argument('--gcp_service_account_path', type=str, default='', help='The project ID in Google Cloud Platform')
parser.add_argument('--datalab_id', type=str, default="'resource_tags_user_user_tag'",
                    help='Column name in report file that contains '
                         'datalab id tag')
parser.add_argument('--usage_date', type=str, default='line_item_usage_start_date',
                    help='Column name in report file that contains '
                         'usage date tag')
parser.add_argument('--product', type=str, default='product_product_name',
                    help='Column name in report file that contains '
                         'product name tag')
parser.add_argument('--usage_type', type=str, default='line_item_usage_type',
                    help='Column name in report file that contains '
                         'usage type tag')
parser.add_argument('--usage', type=str, default='line_item_usage_amount',
                    help='Column name in report file that contains '
                         'usage tag')
parser.add_argument('--cost', type=str, default='line_item_blended_cost',
                    help='Column name in report file that contains cost tag')
parser.add_argument('--resource_id', type=str, default='line_item_resource_id',
                    help='Column name in report file that contains '
                         'datalab resource id tag')
parser.add_argument('--ldap_hostname', type=str, default='localhost', help='Ldap instance hostname')
parser.add_argument('--ldap_dn', type=str, default='dc=example,dc=com',
                    help='Ldap distinguished name')
parser.add_argument('--ldap_ou', type=str, default='ou=People', help='Ldap organisation unit')
parser.add_argument('--ldap_service_username', type=str, default='cn=service-user', help='Ldap service user name')
parser.add_argument('--ldap_service_password', type=str, default='service-user-password',
                    help='Ldap password for admin user')
parser.add_argument('--keycloak_realm_name', type=str, default='datalab', help='Keycloak Realm name')
parser.add_argument('--keycloak_auth_server_url', type=str, default='datalab', help='Keycloak auth server URL')
parser.add_argument('--keycloak_client_name', type=str, default='datalab', help='Keycloak client name')
parser.add_argument('--keycloak_client_secret', type=str, default='datalab', help='Keycloak client secret')
parser.add_argument('--keycloak_user', type=str, default='datalab', help='Keycloak user')
parser.add_argument('--keycloak_user_password', type=str, default='keycloak-user-password',
                    help='Keycloak user password')
parser.add_argument('--tags', type=str, default='line_item_operation,line_item_line_item_description',
                    help='Column name in report file that '
                         'contains tags')
parser.add_argument('--billing_dataset_name', type=str, default='', help='Name of GCP dataset (BigQuery service)'
                                                                         ' for billing')
parser.add_argument('--default_endpoint_name', type=str, default='local', help='Name of localhost provisioning service,'
                                                                               'that created by default')
parser.add_argument('--conf_stepcerts_enabled', type=str, default='false', help='Enable or disable step certificates')
parser.add_argument('--conf_stepcerts_root_ca', type=str, default='', help='Step root CA')
parser.add_argument('--conf_stepcerts_kid', type=str, default='', help='Step KID')
parser.add_argument('--conf_stepcerts_kid_password', type=str, default='', help='Step KID password')
parser.add_argument('--conf_stepcerts_ca_url', type=str, default='', help='Step CA URL')
parser.add_argument('--conf_letsencrypt_enabled', type=str, default='false',
                    help='Enable or disable Let`s Encrypt certificates')
parser.add_argument('--conf_repository_user', type=str, default='',
                    help='user to access repository (used for jars download)')
parser.add_argument('--conf_release_tag', type=str, default='2.5',
                    help='tag used for jars download')
parser.add_argument('--conf_repository_pass', type=str, default='',
                    help='password to access repository (used for jars download)')
parser.add_argument('--conf_repository_address', type=str, default='',
                    help='address to access repository (used for jars download)')
parser.add_argument('--conf_repository_port', type=str, default='',
                    help='port to access repository (used for jars download)')
parser.add_argument('--conf_download_jars', type=str, default='false',
                    help='whether to download jars from repository (True) or to build from sources (False)')
parser.add_argument('--conf_download_docker_images', type=str, default='false',
                    help='whether to download docker images from repository (True) or to build from sources (False)')
parser.add_argument('--conf_letsencrypt_domain_name', type=str, default='', help='Domain names to apply. '
                                                                                 'For multiple domains enter a comma separated list of domains as a parameter'
                                                                                 'ssn.domain_name will be used for ssn_node, DNS A record have to exist during deployment')
parser.add_argument('--conf_letsencrypt_email', type=str, default='', help='Email that will be entered during '
                                                                           'certificate obtaining and can be user for urgent renewal and security notices. '
                                                                           'Use comma to register multiple emails, e.g. u1@example.com,u2@example.com.')
parser.add_argument('--action', required=True, type=str, default='', choices=['build', 'deploy', 'create', 'terminate'],
                    help='Available options: build, deploy, create, terminate')
args = parser.parse_args()


def generate_docker_command():
    docker_command = ''
    command = []
    if args.action == 'terminate':
        command.append('sudo docker run -i ')
    else:
        command.append('sudo docker run -i -v {0}{1}.pem:/root/keys/{1}.pem -v {2}/web_app:/root/web_app '.
                       format(args.key_path, args.conf_key_name, args.workspace_path))
    if args.conf_cloud_provider == 'azure':
        command.append('-v {}:/root/azure_auth.json '.format(args.azure_auth_path))
    elif args.conf_cloud_provider == 'gcp':
        command.append('-v {}:/root/service_account.json '.format(args.gcp_service_account_path))
    if args.ssl_cert_path != '' and args.ssl_key_path != '':
        command.append('-v {}:/root/certs/datalab.crt -v {}:/root/certs/datalab.key '.format(args.ssl_cert_path,
                                                                                             args.ssl_key_path))
    attrs = vars(args)
    skipped_parameters = ['action', 'key_path', 'workspace_path', 'gcp_service_account_path', 'ssl_cert_path',
                          'ssl_key_path']
    for i in attrs:
        if attrs[i] and i not in skipped_parameters:
            command.append("-e '{}={}' ".format(i, attrs[i]))
    command.append('-e "conf_resource=ssn" ')
    command.append('docker.datalab-ssn ')
    command.append('--action {} '.format(args.action))
    return docker_command.join(command)


def build_docker_images(args):
    if args.conf_repository_user and args.conf_repository_pass and args.conf_repository_port and args.conf_repository_address and args.conf_download_docker_images == 'true':
        subprocess.run( 'sudo docker login -u {0} -p {1} {2}:{3}'
                        .format(args.conf_repository_user, args.conf_repository_pass, args.conf_repository_address, args.conf_repository_port), shell=True, check=True)
        subprocess.run('sudo docker pull {}:{}/docker.datalab-base-{}'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
        subprocess.run('sudo docker image tag {}:{}/docker.datalab-base-{} docker.datalab-base'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
        subprocess.run('sudo docker image rm {}:{}/docker.datalab-base-{}'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
        subprocess.run('sudo docker pull {}:{}/docker.datalab-ssn-{}'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
        subprocess.run('sudo docker image tag {}:{}/docker.datalab-ssn-{} docker.datalab-ssn'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
        subprocess.run('sudo docker image rm {}:{}/docker.datalab-ssn-{}'.format(args.conf_repository_address, args.conf_repository_port, args.conf_cloud_provider), shell=True, check=True)
    else:
        # Building base and ssn docker images
        subprocess.run('cd {2}; sudo docker build --build-arg OS={0} --build-arg SRC_PATH="infrastructure-provisioning/src/" --file '
              'infrastructure-provisioning/src/general/files/{1}/'
              'base_Dockerfile -t docker.datalab-base .'.format(args.conf_os_family, args.conf_cloud_provider, args.workspace_path), shell=True, check=True)
        subprocess.run('cd {2}; sudo docker build --build-arg OS={0} --file infrastructure-provisioning/src/general/files/{1}/'
              'ssn_Dockerfile -t docker.datalab-ssn .'.format(args.conf_os_family, args.conf_cloud_provider, args.workspace_path), shell=True, check=True)


def deploy_datalab(args):
    # Creating SSN node
    docker_command = generate_docker_command()
    subprocess.run(docker_command, shell=True, check=True)


def terminate_datalab(args):
    # Dropping datalab environment with selected infrastructure tag
    docker_command = generate_docker_command()
    subprocess.run(docker_command, shell=True, check=True)


if __name__ == "__main__":
    if not args.workspace_path:
        print("Workspace path isn't set, using current directory: {}".format(os.environ['PWD']))
        args.workspace_path = os.environ['PWD']
    if args.action == 'build':
        build_docker_images(args)
    elif args.action == 'deploy':
        deploy_datalab(args)
    elif args.action == 'create':
        build_docker_images(args)
        deploy_datalab(args)
    elif args.action == 'terminate':
        build_docker_images(args)
        terminate_datalab(args)