# Copyright 2015-2019 Lenovo
#
# 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.

# This provides ability to do HTTPS in a manner like ssh host keys for the
# sake of typical internal management devices.  Compatibility back to python
# 2.6 as is found in commonly used enterprise linux distributions.

import base64
import copy
import gzip
import io
import json
import socket
import ssl
import threading
import os

import six

import pyghmi.exceptions as pygexc

try:
    import Cookie
    import httplib
except ImportError:
    import http.client as httplib
    import http.cookies as Cookie


# Used as the separator for form data
BND = b'TbqbLUSn0QFjx9gxiQLtgBK4Zu6ehLqtLs4JOBS50EgxXJ2yoRMhTrmRXxO1lkoAQdZx16'

# We will frequently be dealing with the same data across many instances,
# consolidate forms to single memory location to get benefits..
uploadforms = {}


class FileUploader(threading.Thread):

    def __init__(self, webclient, url, filename, data=None, formname=None,
                 otherfields=(), formwrap=True, excepterror=True):
        self.wc = webclient
        self.url = url
        self.filename = filename
        self.data = data
        self.otherfields = otherfields
        self.formname = formname
        self.rsp = ''
        self.rspstatus = 500
        self.formwrap = formwrap
        self.excepterror = excepterror
        super(FileUploader, self).__init__()
        if not hasattr(self, 'isAlive'):
            self.isAlive = self.is_alive

    def run(self):
        try:
            self.rsp = self.wc.upload(
                self.url, self.filename, self.data, self.formname,
                otherfields=self.otherfields, formwrap=self.formwrap,
                excepterror=self.excepterror)
            self.rspstatus = self.wc.rspstatus
        except Exception:
            try:
                self.rspstatus = self.wc.rspstatus
            except Exception:
                pass
            raise


class FileDownloader(threading.Thread):

    def __init__(self, webclient, url, savefile):
        self.wc = webclient
        self.url = url
        self.savefile = savefile
        self.exc = None
        super(FileDownloader, self).__init__()
        if not hasattr(self, 'isAlive'):
            self.isAlive = self.is_alive

    def run(self):
        try:
            self.wc.download(self.url, self.savefile)
        except Exception as e:
            self.exc = e


def get_upload_form(filename, data, formname, otherfields, boundary=BND):
    if not boundary:
        boundary = base64.urlsafe_b64encode(os.urandom(54))[:66]
    ffilename = filename.split('/')[-1]
    if not formname:
        formname = ffilename
    try:
        return uploadforms[filename]
    except KeyError:
        try:
            data = data.read()
        except AttributeError:
            pass
        form = b''
        for ofield in otherfields:
            tfield = otherfields[ofield]
            xtra=''
            if isinstance(tfield, dict):
                tfield = json.dumps(tfield)
                xtra = '\r\nContent-Type: application/json'
            form += (b'--' + boundary
                     + '\r\nContent-Disposition: form-data; '
                       'name="{0}"{1}\r\n\r\n{2}\r\n'.format(
                           ofield, xtra, tfield).encode('utf-8'))
        form += (b'--' + boundary
                + '\r\nContent-Disposition: form-data; '
                  'name="{0}"; filename="{1}"\r\n'.format(
                      formname, ffilename).encode('utf-8'))
        form += b'Content-Type: application/octet-stream\r\n\r\n' + data
        form += b'\r\n--' + boundary + b'--\r\n'
        uploadforms[filename] = form, boundary
        return uploadforms[filename]


class SecureHTTPConnection(httplib.HTTPConnection, object):
    default_port = httplib.HTTPS_PORT

    def __init__(self, host, port=None, key_file=None, cert_file=None,
                 ca_certs=None, strict=None, verifycallback=None, clone=None,
                 **kwargs):
        if 'timeout' not in kwargs:
            kwargs['timeout'] = 60
        self.mytimeout = kwargs['timeout']
        self._currdl = None
        self.lastjsonerror = None
        self.broken = False
        self.thehost = host
        self.theport = port
        self._upbuffer = None
        try:
            httplib.HTTPConnection.__init__(self, host, port, strict=strict,
                                            **kwargs)
        except TypeError:
            httplib.HTTPConnection.__init__(self, host, port, **kwargs)
        if clone:
            self._certverify = clone._certverify
            self.cookies = clone.cookies
            self.stdheaders = copy.deepcopy(clone.stdheaders)
        else:
            self._certverify = verifycallback
            self.cookies = {}
            self.stdheaders = {}
        if self._certverify:
            self.cert_reqs = ssl.CERT_NONE  # use custom validation
        else:
            self.cert_reqs = ssl.CERT_REQUIRED  # use standard validation
        if '[' not in host and '%' in host and 'Host'not in self.stdheaders:
            self.stdheaders['Host'] = '[' + host[:host.find('%')] + ']'

    def __del__(self):
        if self.sock:
            self.sock.close()
            self.sock = None

    def dupe(self, timeout=None):
        if timeout is None:
            timeout = self.mytimeout
        return SecureHTTPConnection(self.thehost, self.theport, clone=self,
                                    timeout=timeout)

    def set_header(self, key, value):
        self.stdheaders[key] = value

    def set_basic_credentials(self, username, password):
        if isinstance(username, bytes) and not isinstance(username, str):
            username = username.decode('utf-8')
        if isinstance(password, bytes) and not isinstance(password, str):
            password = password.decode('utf-8')
        authinfo = ':'.join((username, password))
        if not isinstance(authinfo, bytes):
            authinfo = authinfo.encode('utf-8')
        authinfo = base64.b64encode(authinfo)
        if not isinstance(authinfo, str):
            authinfo = authinfo.decode('utf-8')
        self.stdheaders['Authorization'] = 'Basic {0}'.format(authinfo)

    def connect(self):
        addrinfo = socket.getaddrinfo(self.host, self.port)[0]
        # workaround problems of too large mtu, moderately frequent occurance
        # in this space
        plainsock = socket.socket(addrinfo[0])
        plainsock.settimeout(self.mytimeout)
        try:
            plainsock.setsockopt(socket.IPPROTO_TCP, socket.TCP_MAXSEG, 1456)
        except socket.error:
            pass
        plainsock.connect(addrinfo[4])
        if self._certverify:
            ctx = ssl.create_default_context()
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            self.sock = ctx.wrap_socket(plainsock)
            bincert = self.sock.getpeercert(binary_form=True)
            if not self._certverify(bincert):
                raise pygexc.UnrecognizedCertificate('Unknown certificate',
                                                     bincert)
        else:
            ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
            ctx.load_default_certs()
            ctx.verify_mode = ssl.CERT_REQUIRED
            ctx.check_hostname = True
            self.sock = ctx.wrap_socket(plainsock,
                                        server_hostname=self.thehost)

    def getresponse(self):
        try:
            rsp = super(SecureHTTPConnection, self).getresponse()
            try:
                hdrs = [x.split(':', 1) for x in rsp.msg.headers]
            except AttributeError:
                hdrs = rsp.msg.items()
            for hdr in hdrs:
                if hdr[0] == 'Set-Cookie':
                    c = Cookie.BaseCookie(hdr[1])
                    for k in c:
                        self.cookies[k] = c[k].value
        except httplib.BadStatusLine:
            self.broken = True
            raise
        return rsp

    def grab_json_response(self, url, data=None, referer=None, headers=None):
        self.lastjsonerror = None
        body, status = self.grab_json_response_with_status(
            url, data, referer, headers)
        if status == 200:
            return body
        self.lastjsonerror = body
        return {}

    def grab_json_response_with_status(self, url, data=None, referer=None,
                                       headers=None, method=None):
        webclient = self.dupe()
        if isinstance(data, dict):
            data = json.dumps(data)
        if data:
            if not method:
                method = 'POST'
            webclient.request(method, url, data, referer=referer,
                              headers=headers)
        else:
            if not method:
                method = 'GET'
            webclient.request(method, url, referer=referer, headers=headers)
        try:
            rsp = webclient.getresponse()
        except httplib.BadStatusLine:
            return 'Target Unavailable', 500
        except ssl.SSLError as e:
            if 'timed out' in str(e):
                return 'Target Unavailable', 500
            raise
        body = rsp.read()
        if rsp.getheader('Content-Encoding', None) == 'gzip':
            try:
                body = gzip.GzipFile(fileobj=io.BytesIO(body)).read()
            except IOError:
                # some implementations will send non-gzipped and claim it as
                # gzip
                pass
        if rsp.status >= 200 and rsp.status < 300:
            if body and not isinstance(body, type(u'')):
                try:
                    body = body.decode('utf8')
                except Exception:
                    body = body.decode('iso-8859-1')
            return json.loads(body) if body else {}, rsp.status
        return body, rsp.status

    def grab_rsp(self, url, data=None, referer=None, headers=None, method=None):
        webclient = self.dupe()
        if isinstance(data, dict):
            data = json.dumps(data)
        if data:
            if not method:
                method = 'POST'
            webclient.request(method, url, data, referer=referer,
                              headers=headers)
        else:
            if not method:
                method = 'GET'
            webclient.request(method, url, referer=referer, headers=headers)
        rsp = webclient.getresponse()
        return rsp

    def download(self, url, file):
        """Download a file to filename or file object

        """
        if isinstance(file, six.string_types):
            file = open(file, 'wb')
        webclient = self.dupe()
        dlheaders = self.stdheaders.copy()
        if 'Accept-Encoding' in dlheaders:
            del dlheaders['Accept-Encoding']
        webclient.request('GET', url, headers=dlheaders)
        rsp = webclient.getresponse()
        self._currdl = rsp
        self._dlfile = file
        chunk = rsp.read(16384)
        while chunk:
            file.write(chunk)
            chunk = rsp.read(16384)
        self._currdl = None
        file.close()

    def get_download_progress(self):
        if not self._currdl:
            return None
        return float(self._dlfile.tell()) / float(
            self._currdl.getheader('content-length'))

    def upload(self, url, filename, data=None, formname=None,
               otherfields=(), formwrap=True, excepterror=True):
        """Upload a file to the url

        :param url:
        :param filename: The name of the file
        :param data: A file object or data to use rather than reading from
                     the file.
        :return:
        """
        if data is None:
            data = open(filename, 'rb')
        ulhdrs = self.stdheaders.copy()
        if formwrap:
            guf = get_upload_form(
                filename, data, formname, otherfields, boundary=None)
            self._upbuffer = io.BytesIO(guf[0])
            ulhdrs['Content-Type'] = b'multipart/form-data; boundary=' + guf[1]
            ulhdrs['Content-Length'] = len(uploadforms[filename][0])
            self.ulsize = len(uploadforms[filename][0])
        else:
            canseek = True
            try:
                curroff = data.tell()
            except io.UnsupportedOperation:
                canseek = False
                databytes = data.read()
                self.ulsize = len(databytes)
                self._upbuffer = io.BytesIO(databytes)
            if canseek:
                data.seek(0, 2)
                self.ulsize = data.tell()
                data.seek(curroff, 0)
                self._upbuffer = data
            ulhdrs['Content-Type'] = b'application/octet-stream'
            ulhdrs['Content-Length'] = self.ulsize
        webclient = self.dupe()
        webclient.request('POST', url, self._upbuffer, ulhdrs)
        rsp = webclient.getresponse()
        # peer updates in progress should already have pointers,
        # subsequent transactions will cause memory to needlessly double,
        # but easiest way to keep memory relatively low
        if formwrap:
            try:
                del uploadforms[filename]
            except KeyError:  # something could have already deleted it
                pass
        self.rspstatus = rsp.status
        if excepterror and (rsp.status < 200 or rsp.status >= 300):
            raise Exception('Unexpected response in file upload: %s'
                            % rsp.read())
        body = rsp.read()
        if rsp.getheader('Content-Encoding', None) == 'gzip':
            try:
                body = gzip.GzipFile(fileobj=io.BytesIO(body)).read()
            except IOError:
                # In case the implementation lied, let the body return
                # unprocessed
                pass
        return body

    def get_upload_progress(self):
        if self._upbuffer is None:
            return 0.0
        return float(self._upbuffer.tell()) / float(self.ulsize)

    def request(self, method, url, body=None, headers=None, referer=None):
        if headers is None:
            headers = self.stdheaders.copy()
        else:
            headers = headers.copy()
        if method == 'GET' and 'Content-Type' in headers:
            del headers['Content-Type']
        if method == 'POST' and body and 'Content-Type' not in headers:
            headers['Content-Type'] = 'application/x-www-form-urlencoded'
        if body and 'Content-Length' not in headers:
            headers['Content-Length'] = len(body)
        if self.cookies:
            cookies = []
            for ckey in self.cookies:
                cookies.append('{0}={1}'.format(ckey, self.cookies[ckey]))
            cookies_header = '; '.join(cookies)
            if headers.get('Cookie', None) is None:
                headers['Cookie'] = cookies_header
            else:
                headers['Cookie'] += '; ' + '; '.join(cookies)
        if referer:
            headers['Referer'] = referer
        try:
            return super(SecureHTTPConnection, self).request(method, url, body,
                                                             headers)
        except httplib.CannotSendRequest:
            self.broken = True
            raise

