code/default/lib/noarch/front_base/pyopenssl_wrap.py - master - githubmirror/XX-Net

code/default/lib/noarch/front_base/pyopenssl_wrap.py
535 строк18 KB
Micheal X
5.8.3 fix http1 bug.
25 окт 2023, 07:28
25 окт 2023, 07:28f608910

# this wrap has a close callback.
# Which is used by  ip manager
#  ip manager keep a connection number counter for every ip.

# the wrap SSL implementation, python 2.7 will use pyOpenSSL, python 3.x will use build in ssl.
# This can also be used to store some attribute like ip_str/appid

import os
import sys
import datetime
import socket
import ssl
import OpenSSL
import time
import select
import errno

import utils


class SSLConnection(object):
    def __init__(self, context, sock, ip_str=None, sni=None, on_close=None):
        self._context = context
        self._sock = sock
        self.ip_str = ip_str
        self.sni = sni
        self._makefile_refs = 0
        self._on_close = on_close
        self.peer_cert = None
        self.socket_closed = False
        self.timeout = self._sock.gettimeout() or 0.1
        self.running = True
        self._connection = None
        self.wrap()

    def wrap(self):
        ip, port = utils.get_ip_port(self.ip_str)
        if isinstance(ip, str):
            ip = utils.to_bytes(ip)

        ip_port = (utils.to_str(ip), port)

        if sys.version_info[0] == 3:
            try:
                self._sock.connect((ip, port))
            except Exception as e:
                raise socket.error('conn %s fail, sni:%s, e:%r' % (self.ip_str, self.sni, e))

            self._connection = self._context.wrap_socket(self._sock, server_hostname=self.sni,
                                                         do_handshake_on_connect=False)
        else:
            self._connection = OpenSSL.SSL.Connection(self._context, self._sock)
            self._connection.set_connect_state()
            if self.sni:
                try:
                    self._connection.set_tlsext_host_name(self.sni)
                except:
                    pass
            self._connection.connect(ip_port)

    def is_support_h2(self):
        if sys.version_info[0] == 3:
            return self._connection.selected_alpn_protocol() == "h2" or self._connection.selected_npn_protocol() == "h2"
        else:
            return self._connection.get_alpn_proto_negotiated()

    def __getattr__(self, attr):
        if attr == "socket_closed":
            # work around in case close before finished init.
            return True

        elif attr in ('is_support_h2', "_on_close", '_context', '_sock', '_connection', '_makefile_refs',
                      'sni', 'wrap', 'socket_closed'):
            return getattr(self, attr)

        elif hasattr(self._connection, attr):
            return getattr(self._connection, attr)

    def __del__(self):
        if not self.socket_closed and self._connection:
            self._connection.close()
            self.socket_closed = True

        if self._on_close:
            self._on_close(self.ip_str, self.sni)
            self._on_close = None

    def set_tlsext_host_name(self, hostname):
        self._connection.server_hostname = utils.to_str(hostname)

    def get_cert(self):
        # py3 only
        if self.peer_cert:
            return self.peer_cert

        cert = self._connection.getpeercert()
        # For debug:
        #cert = self._connection.getpeercert(True)
        #from asn1crypto.x509 import Certificate
        #cert = Certificate.load(cert)

        self.peer_cert = {
            "cert": cert,
            "issuer_commonname": "",
            "commonName": "",
            "altName": []
        }
        for kv in cert["issuer"]:
            k, v = kv[0]
            if k == 'commonName':
                self.peer_cert["issuer_commonname"] = v

        for kv in cert["subject"]:
            k, v = kv[0]
            if k == 'commonName':
                self.peer_cert["commonName"] = v

        for k, v in cert["subjectAltName"]:
            self.peer_cert["altName"].append(v)
        self.peer_cert["altName"] = tuple(self.peer_cert["altName"])

        return self.peer_cert

    def __iowait(self, io_func, *args, **kwargs):
        fd = self._sock.fileno()
        time_start = time.time()
        while self.running:
            time_now = time.time()
            wait_timeout = max(0.1, self.timeout - (time_now - time_start))
            wait_timeout = min(wait_timeout, 10)
            # in case socket was blocked by FW
            # recv is called before send request, which timeout is 240
            # then send request is called and timeout change to 100

            try:
                return io_func(*args, **kwargs)
            except (OpenSSL.SSL.WantReadError, OpenSSL.SSL.WantX509LookupError) as e:
                sys.exc_clear()
                _, _, errors = select.select([fd], [], [fd], wait_timeout)
                if errors:
                    raise e
                if time_now - time_start > self.timeout:
                    break
            except OpenSSL.SSL.WantWriteError as e:
                sys.exc_clear()
                _, _, errors = select.select([], [fd], [fd], wait_timeout)
                if errors:
                    raise e
                time_now = time.time()
                if time_now - time_start > self.timeout:
                    break
            except OpenSSL.SSL.SysCallError as e:
                if e[0] == 10035 and 'WSAEWOULDBLOCK' in e[1]:
                    sys.exc_clear()
                    if io_func == self._connection.send:
                        _, _, errors = select.select([], [fd], [fd], wait_timeout)
                    else:
                        _, _, errors = select.select([fd], [], [fd], wait_timeout)

                    if errors:
                        raise e
                    time_now = time.time()
                    if time_now - time_start > self.timeout:
                        break
                else:
                    raise e
            except Exception as e:
                #self.logger.exception("e:%r", e)
                raise e

        return 0

    def accept(self):
        sock, addr = self._sock.accept()
        client = OpenSSL.SSL.Connection(sock._context, sock)
        return client, addr

    def do_handshake(self):
        self.__iowait(self._connection.do_handshake)

    def connect(self, *args, **kwargs):
        return self.__iowait(self._connection.connect, *args, **kwargs)

    def __send(self, data, flags=0):
        try:
            return self.__iowait(self._connection.send, data, flags)
        except OpenSSL.SSL.SysCallError as e:
            if e[0] == -1 and not data:
                # errors when writing empty strings are expected and can be ignored
                return 0
            raise
        except Exception as e:
            #self.logger.exception("ssl send:%r", e)
            raise

    def __send_memoryview(self, data, flags=0):
        if hasattr(data, 'tobytes'):
            data = data.tobytes()
        return self.__send(data, flags)

    send = __send if sys.version_info >= (2, 7, 5) else __send_memoryview

    def recv(self, bufsiz, flags=0):
        pending = self._connection.pending()
        if pending:
            return self._connection.recv(min(pending, bufsiz))

        try:
            return self.__iowait(self._connection.recv, bufsiz, flags)
        except OpenSSL.SSL.ZeroReturnError:
            return ''
        except OpenSSL.SSL.SysCallError as e:
            if e[0] == -1 and 'Unexpected EOF' in e[1]:
                # remote closed
                #raise e
                return ""
            elif e[0] == 10053 or e[0] == 10054 or e[0] == 10038:
                return ""
            raise

    def recv_into(self, buf, nbytes=None):
        pending = self._connection.pending()
        if pending:
            ret = self._connection.recv_into(buf, nbytes)
            if not ret:
                # self.logger.debug("recv_into 0")
                pass
            return ret

        while self.running:
            try:
                ret = self.__iowait(self._connection.recv_into, buf, nbytes)
                if not ret:
                    # self.logger.debug("recv_into 0")
                    pass
                return ret
            except OpenSSL.SSL.ZeroReturnError as e:
                raise e
            except OpenSSL.SSL.SysCallError as e:
                if sys.version_info[0] == 2:
                    if e[0] == -1 and 'Unexpected EOF' in e[1]:
                        # errors when reading empty strings are expected and can be ignored
                        return 0
                    elif e[0] == 11 and e[1] == 'EAGAIN':
                        continue
                raise
            except Exception as e:
                if sys.version_info[0] == 2 and e == errno.EAGAIN:
                    continue
                # logging.exception("recv %r", e)
                raise e

    def read(self, bufsiz, flags=0):
        return self.recv(bufsiz, flags)

    def write(self, buf, flags=0):
        return self.sendall(buf, flags)

    def close(self, reason=""):
        if self._makefile_refs < 1:
            self.running = False
            if not self.socket_closed:
                socket.socket.close(self._sock)
                self.socket_closed = True
                if self._on_close:
                    self._on_close(self.ip_str, self.sni, reason=reason)
                    self._on_close = None
        else:
            self._makefile_refs -= 1

    def settimeout(self, t):
        if not self.running:
            return

        if self.timeout != t:
            if sys.version_info[0] == 3:
                self._connection.settimeout(t)
            else:
                self._sock.settimeout(t)
            self.timeout = t

    def makefile(self, mode='r', bufsize=-1):
        self._makefile_refs += 1
        return socket._fileobject(self, mode, bufsize, close=True)


class SSLContext(object):
    def __init__(self, logger, ca_certs=None, cipher_suites=None, support_http2=True, protocol=None):
        self.logger = logger

        if sys.version_info[0] == 3:
            if protocol == "TLSv1_2":
                ssl_version = ssl.PROTOCOL_TLSv1_2

            elif hasattr(ssl, "PROTOCOL_TLS"):
                ssl_version = ssl.PROTOCOL_TLS
            elif hasattr(ssl, "PROTOCOL_TLSv1_2"):
                ssl_version = ssl.PROTOCOL_TLSv1_2
            elif hasattr(ssl, "PROTOCOL_TLSv1_1"):
                ssl_version = ssl.PROTOCOL_TLSv1_1
            elif hasattr(ssl, "PROTOCOL_TLSv1"):
                ssl_version = ssl.PROTOCOL_TLSv1
            elif hasattr(ssl, "PROTOCOL_SSLv3"):
                ssl_version = ssl.PROTOCOL_SSLv3
            elif hasattr(ssl, "PROTOCOL_SSLv2"):
                ssl_version = ssl.PROTOCOL_SSLv2
            else:
                ssl_version = ssl.PROTOCOL_SSLv23

            self.logger.info("SSL use version:%s", self.supported_protocol())
            self.context = ssl.SSLContext(protocol=ssl_version)

            self.set_ca(ca_certs)

            if cipher_suites:
                self.context.set_ciphers(':'.join(cipher_suites))

            self.support_alpn_npn = None
            if support_http2:
                try:
                    self.context.set_alpn_protocols(['h2', 'http/1.1'])
                    self.logger.info("OpenSSL support alpn")
                    self.support_alpn_npn = "alpn"
                    return
                except Exception as e:
                    self.logger.exception("set_alpn_protos:%r", e)
                    pass

                try:
                    self.context.set_npn_protocols(['h2', 'http/1.1'])
                    self.logger.info("OpenSSL support npn")
                    self.support_alpn_npn = "npn"
                except Exception as e:
                    # xlog.exception("set_npn_select_callback:%r", e)
                    self.logger.info("OpenSSL dont't support npn/alpn, no HTTP/2 supported.")
                    pass
        else:
            self.ssl_version = self.supported_protocol()

            protocol_version = getattr(OpenSSL.SSL, '%s_METHOD' % self.ssl_version)
            self.context = OpenSSL.SSL.Context(protocol_version)

            self.set_ca(ca_certs)

            if cipher_suites:
                self.context.set_cipher_list(':'.join(cipher_suites))

            self.support_alpn_npn = None
            if support_http2:
                try:
                    self.context.set_alpn_protos([b'h2', b'http/1.1'])
                    self.logger.info("OpenSSL support alpn")
                    self.support_alpn_npn = "alpn"
                    return
                except Exception as e:
                    # self.logger.exception("set_alpn_protos:%r", e)
                    pass

                try:
                    self.context.set_npn_select_callback(SSLContext.npn_select_callback)
                    self.logger.info("OpenSSL support npn")
                    self.support_alpn_npn = "npn"
                except Exception as e:
                    #xlog.exception("set_npn_select_callback:%r", e)
                    self.logger.info("OpenSSL dont't support npn/alpn, no HTTP/2 supported.")
                    pass

    @staticmethod
    def npn_select_callback(conn, protocols):
        # self.logger.debug("npn protocl:%s", ";".join(protocols))
        if b"h2" in protocols:
            conn.protos = "h2"
            return b"h2"
        else:
            return b"http/1.1"

    @staticmethod
    def supported_protocol():
        if sys.version_info[0] == 3:
            if hasattr(ssl, "HAS_TLSv1_3") and ssl.HAS_TLSv1_3:
                ssl_version = "TLSv1_3"
            elif hasattr(ssl, "HAS_TLSv1_2") and ssl.HAS_TLSv1_2:
                ssl_version = "TLSv1_2"
            elif hasattr(ssl, "HAS_TLSv1_1") and ssl.HAS_TLSv1_1:
                ssl_version = "TLSv1_1"
            elif hasattr(ssl, "HAS_SSLv3") and ssl.HAS_SSLv3:
                ssl_version = "SSLv3"
            elif hasattr(ssl, "HAS_SSLv2") and ssl.HAS_SSLv2:
                ssl_version = "SSLv2"
            else:
                ssl_version = "SSLv1"
        else:
            if hasattr(OpenSSL.SSL, "TLSv1_2_METHOD"):
                ssl_version = "TLSv1_2"
            elif hasattr(OpenSSL.SSL, "TLSv1_1_METHOD"):
                ssl_version = "TLSv1_1"
            elif hasattr(OpenSSL.SSL, "TLSv1_METHOD"):
                ssl_version = "TLSv1"
            else:
                ssl_version = "SSLv23"

            if sys.platform == "darwin":
                # MacOS pyOpenSSL has TLSv1_2_METHOD attr but can use.
                # There for we hard code here.
                # may be try/cache is a better solution.
                ssl_version = "TLSv1"

            # freenas openssl support fix from twitter user "himanzero"
            # https://twitter.com/himanzero/status/645231724318748672
            if sys.platform == "freebsd9":
                ssl_version = "TLSv1"

        return ssl_version

    def set_ca(self, ca_certs):
        if sys.version_info[0] == 3:
            try:
                if ca_certs:
                    self.context.load_verify_locations(cafile=os.path.abspath(ca_certs))
                    self.context.verify_mode = ssl.CERT_REQUIRED
                else:
                    self.context.verify_mode = ssl.CERT_NONE
            except Exception as e:
                self.logger.debug("set_ca fail:%r", e)
        else:
            if ca_certs:
                self.context.load_verify_locations(os.path.abspath(ca_certs))
                self.context.set_verify(OpenSSL.SSL.VERIFY_PEER, lambda c, x, e, d, ok: ok)
            else:
                self.context.set_verify(OpenSSL.SSL.VERIFY_NONE, lambda c, x, e, d, ok: ok)


from pyasn1.codec.der.decoder import decode
from pyasn1.error import PyAsn1Error
from pyasn1.type import univ, constraint, char, namedtype, tag


class _GeneralName(univ.Choice):
    # We are only interested in dNSNames. We use a default handler to ignore
    # other types.
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('dNSName', char.IA5String().subtype(
                implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 2)
            )
        ),
    )


class _GeneralNames(univ.SequenceOf):
    componentType = _GeneralName()
    sizeSpec = univ.SequenceOf.sizeSpec + constraint.ValueSizeConstraint(1, 1024)


class SSLCert:
    def __init__(self, cert):
        """
            Returns a (common name, [subject alternative names]) tuple.
        """
        self.x509 = cert

    @classmethod
    def from_pem(klass, txt):
        x509 = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, txt)
        return klass(x509)

    @classmethod
    def from_der(klass, der):
        pem = ssl.DER_cert_to_PEM_cert(der)
        return klass.from_pem(pem)

    def to_pem(self):
        return OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, self.x509)

    def digest(self, name):
        return self.x509.digest(name)

    @property
    def issuer(self):
        return self.x509.get_issuer().get_components()

    @property
    def notbefore(self):
        t = self.x509.get_notBefore()
        return datetime.datetime.strptime(t, "%Y%m%d%H%M%SZ")

    @property
    def notafter(self):
        t = self.x509.get_notAfter()
        return datetime.datetime.strptime(t, "%Y%m%d%H%M%SZ")

    @property
    def has_expired(self):
        return self.x509.has_expired()

    @property
    def subject(self):
        return self.x509.get_subject().get_components()

    @property
    def serial(self):
        return self.x509.get_serial_number()

    @property
    def keyinfo(self):
        pk = self.x509.get_pubkey()
        types = {
            OpenSSL.crypto.TYPE_RSA: "RSA",
            OpenSSL.crypto.TYPE_DSA: "DSA",
        }
        return (
            types.get(pk.type(), "UNKNOWN"),
            pk.bits()
        )

    @property
    def cn(self):
        c = None
        for i in self.subject:
            if i[0] == "CN":
                c = i[1]
        return c

    @property
    def altnames(self):
        altnames = []
        for i in range(self.x509.get_extension_count()):
            ext = self.x509.get_extension(i)
            if ext.get_short_name() == "subjectAltName":
                try:
                    dec = decode(ext.get_data(), asn1Spec=_GeneralNames())
                except PyAsn1Error:
                    continue
                for i in dec[0]:
                    altnames.append(i[0].asOctets())
        return altnames