# -*-coding:UTF-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2022-2023. All rights reserved.
import stat
import json
import os
import requests
from requests.exceptions import SSLError
from requests_toolbelt import MultipartEncoder
from utils.common.ssh_util2 import Ssh
import utils.common.log as logger


class RestClient(object):
    def __init__(self, ip, port):
        self.ip = ip
        self.port = port
        self.token = None
        self.jsession_id = None
        self.https = "https://"
        self.login_url = '/api/v2/aa/sessions'
        self.cert = False

    @staticmethod
    def format_data(data):
        if data is not None:
            json_data = json.dumps(data)
        else:
            json_data = data
        return json_data

    @staticmethod
    def login_with_cert(url, json_data, login_header, session, cert_list):
        for cert in cert_list:
            cert_path = os.path.join(cert[0], cert[1])
            try:
                res = session.post(url, data=json_data, headers=login_header, verify=cert_path)
            except SSLError:
                logger.error("Certificate[%s] verification failed." % cert[1])
                continue
            cert = cert_path
            res.close()
            return res, cert

        raise Exception("No certificate available")

    def make_header(self, content_type='application/json', **kwargs):
        """
        构造header
        :param content_type:
        :param token:
        :param kwargs:
        :return:
        """
        header = {}
        if content_type is not None:
            header = {'Content-type': content_type,
                      'User-Agent': 'Mozilla/5.0'}
        if self.token is not None:
            header['X-Auth-Token'] = self.token
        for i in kwargs:
            header[i] = kwargs.get(i)
        return header

    def set_cookie(self, token):
        """
        设置token
        :param token:
        :return:
        """
        self.token = token

    def get_cert(self):
        """
        判断当前request请求是否需要证书
        :param:
        :return bool:
        """
        cert_path = os.path.join(__file__, '../../cert')
        if os.path.exists(cert_path):
            ip_cert_path = os.path.join(cert_path, self.ip)
            cert_list = os.listdir(ip_cert_path)
            if cert_list:
                certs = [(ip_cert_path, cert) for cert in cert_list]
                return certs
            else:
                err_msg = "Failed to obtain the certificate."
                logger.error(err_msg)
                raise Exception(err_msg)
        else:
            return False

    def login(self, username, password):
        """
        登陆模块
        :param username:用户名
        :param password:密码
        :return:
        """
        url = self.https + self.ip + ":" + self.port + self.login_url
        logger.info("login request: %s %s " % ("post", url))
        login_data = {"username": username,
                      "password": password,
                      "scope": 0}
        json_data = json.dumps(login_data)
        login_header = {
            'Content-type': 'application/json',
            'Cookie': '__LANGUAGE_KEY__=zh-CN; __IBASE_LANGUAGE_KEY__=zh-CN'}
        cert_list = self.get_cert()
        with requests.session() as session:
            if cert_list:
                res, cert = self.login_with_cert(url, json_data, login_header, session, cert_list)
                self.cert = cert
            else:
                requests.packages.urllib3.disable_warnings()
                res = session.post(url, data=json_data, headers=login_header, verify=self.cert)
                res.close()
        return res

    def normal_request(self, url, method, data=None):
        """
        一般的请求，除登陆和上传文件外的请求统一走这里
        :param url:
        :param method:
        :return:
        """
        logger.info("request: %s %s " % (method, url))
        headers = self.make_header()
        # 消除规避ssl认证导致的警告
        requests.packages.urllib3.disable_warnings()
        json_data = self.format_data(data)
        with requests.session() as session:
            res = self.session_request(*[headers, json_data, method, session, url])
            res.close()
        return res

    def session_request(self, *args):
        headers, json_data, method, session, url = args
        if method == 'put':
            res = session.put(url, data=json_data, headers=headers,
                              verify=self.cert)
        elif method == 'post':
            res = session.post(url, data=json_data, headers=headers,
                               verify=self.cert)
        elif method == 'get':
            res = session.get(url, headers=headers, verify=self.cert)
        elif method == 'delete':
            res = session.delete(url, headers=headers, verify=self.cert)
        return res

    def make_upload_header(self, content_type='application/json', **kwargs):
        """
        构造header
        :param content_type:
        :param token:
        :param kwargs:
        :return:
        """
        header = {}
        if content_type is not None:
            header = {'Content-type': content_type,
                      'User-Agent': 'Mozilla/5.0'}
        if self.token is not None:
            header['X-Auth-Token'] = self.token
        for i in kwargs:
            header[i] = kwargs.get(i)
        return header

    def upload_request(self, url, method, data=None):
        """
        一般的请求，除登陆和上传文件外的请求统一走这里
        :param url:
        :param method:
        :param data:
        :return:
        """
        logger.info("upload request: %s %s " % (method, url))
        content_type = "application/gzip"
        headers = self.make_upload_header(content_type)
        # 消除规避ssl认证导致的警告
        requests.packages.urllib3.disable_warnings()
        json_data = self.format_data(data)
        with requests.session() as session:
            res = self.session_request(*[headers, json_data, method, session, url])
            res.close()
        return res

    def upload_tar_request(self, url, method, file_path):
        """
        上传tar包
        :param file_path:
        :return:
        """
        logger.info("upload request: %s %s " % (method, url))
        requests.packages.urllib3.disable_warnings()
        file_size = os.path.getsize(file_path)
        file_name = os.path.basename(file_path)
        logger.info("file %s size %s" % (file_name, file_size))
        with open(file_path, "rb") as file_object:
            files = {"file": file_object}
            file_data = {"name": file_path,
                         "Content-Type": "application/gzip"
                         }
            headers = self.make_header(content_type=None)
            res = requests.post(
                url=url,
                files=files,
                data=file_data,
                headers=headers,
                verify=self.cert)
            res.close()
        return res

    def upload_tar_request_stream(self, url, method, file_path):
        """
        上传tar包
        :param file_path:
        :return:
        """
        logger.info("upload request: %s %s " % (method, url))
        requests.packages.urllib3.disable_warnings()

        m = MultipartEncoder(fields={'file': (os.path.basename(file_path),
                                              open(file_path, "rb"))})
        headers = self.make_header(content_type=m.content_type)
        res = requests.post(
            url=url,
            data=m,
            headers=headers,
            verify=self.cert)
        res.close()
        return res

    def upload_tar_piece_request(self, url, method, file_path):
        """
        切片上传tar包
        :param url:
        :param method
        :return:
        """
        logger.info("upload request: %s %s " % (method, url))
        chunkfilesize = 10 * 1024 * 1024
        requests.packages.urllib3.disable_warnings()
        file_size = os.path.getsize(file_path)
        md5 = str(int(os.path.getmtime(file_path))) + \
            '_' + str(os.path.getsize(file_path))
        file_name = os.path.basename(file_path)

        # 计算得到切片的数目
        try:
            chunks = (file_size + chunkfilesize - 1) / chunkfilesize
        except ZeroDivisionError as e:
            err_msg = 'Zero Division Error. Detail: %s' % str(e)
            logger.error(err_msg)
            raise e
        chunk_data = chunkfilesize
        with open(file_path, "r") as file_object:
            for chunk in range(0, chunks):
                files = {
                    "tiFile": (
                        "blob",
                        file_object.read(chunk_data),
                        "application/octet-stream")}
                file_data = {"onlySession": md5,
                             "fileName": file_name,
                             "chunks": chunks,
                             "currentChunk": chunk,
                             "type": "application/gzip"
                             }
                headers = self.make_header(content_type=None)
                res = requests.post(
                    url=url,
                    files=files,
                    data=file_data,
                    headers=headers,
                    verify=self.cert)
                res.close()
        return res

    def download_request(self, url, filename, method="get", data=None):
        """
        一般的请求，除登陆和上传文件外的请求统一走这里
        :param url:
        :param method:
        :param filename
        :param data
        :return:
        """
        logger.info("download request: %s %s " % (method, url))
        headers = self.make_header()
        # 消除规避ssl认证导致的警告
        requests.packages.urllib3.disable_warnings()
        json_data = self.format_data(data)
        with requests.session() as session:
            res = self.session_request(*[headers, json_data, method, session, url])
            flag = os.O_CREAT | os.O_RDWR
            mode = stat.S_IRUSR | stat.S_IWUSR
            with os.fdopen(os.open(filename, flag, mode), 'wb') as code:
                code.write(res.content)
            res.close()
        return res


class StorageSSHClient:
    def __init__(self, host_ip, username, password, root_pwd=None):
        self.host_ip = host_ip
        self.username = username
        self.password = password
        self.ssh_client = Ssh.ssh_create_client(host_ip, username, password)
        if username == "root":
            self.is_root = True
        else:
            self.is_root = False
        self.root_pwd = root_pwd

    def __del__(self):
        if self.ssh_client:
            Ssh.ssh_close(self.ssh_client)

    def switch_root(self, root_password=None):
        cmd_res = None
        try:
            if self.is_root is False:
                if root_password is None:
                    root_password = self.root_pwd
                cmd_res = Ssh.ssh_send_command(
                    self.ssh_client, 'su', 'assword:', 20)
                cmd_res = Ssh.ssh_send_command(
                    self.ssh_client, root_password, '#', 20)
                cmd_res = Ssh.ssh_send_command(
                    self.ssh_client, 'TMOUT=0', '#', 20)
                self.is_root = True
        except Exception as e:
            logger.error("Failed to su root on host cmd_res:%s, using user "
                         "%s, manager_ip:%s, err:%s"
                         % (str(cmd_res), self.username, self.host_ip, str(e)))
            raise

    def send_cmd(self, cmd, expect, timeout=60, **kwargs):
        sensitive = kwargs.get("sensitive", False)
        retry_times = kwargs.get("retry_times", 0)
        if not sensitive:
            logger.info("exec cmd: %s, expect: %s" % (cmd, expect))
        else:
            logger.info("exec cmd: %s, expect: %s" % ("****", expect))
        cmd_ret = Ssh.ssh_send_command(
            self.ssh_client, cmd, expect, timeout, retry_times)
        logger.debug("cmd ret: %s" % cmd_ret)
        return cmd_ret

    def send_cmd_return(self, cmd, timeout=60, retry_times=0, sensitive=False):
        if not sensitive:
            logger.info("exec cmd: %s" % cmd)
        else:
            logger.info("exec cmd: %s" % "****")
        return Ssh.ssh_exec_command_return(
            self.ssh_client, cmd, timeout, retry_times)

    def upload(self, local_file, remote_file):
        logger.info("upload file: %s -> %s" % (local_file, remote_file))
        return Ssh.put_file(self.host_ip, self.username, self.password,
                            local_file, remote_file)
