# SPDX-License-Identifier: BSD-3-Clause

"""
The core for a "virtual network" proxy.
"""

# This file is part of Haketilo.
#
# Copyright (c) 2015, inaz2
# Copyright (C) 2021 jahoti <jahoti@tilde.team>
# Copyright (C) 2021 Wojtek Kosior <koszko@koszko.org>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of proxy2 nor the names of its contributors may be used to
#   endorse or promote products derived from this software without specific
#   prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
# I, Wojtek Kosior, thereby promise not to sue for violation of this file's
# license. Although I request that you do not make use this code in a way
# incompliant with the license, I am not going to enforce this in court.

from pathlib import Path
import socket, ssl, subprocess, sys, threading
from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn

lock = threading.Lock()

class ProxyRequestHandler(BaseHTTPRequestHandler):
    """
    Handles a network request made to the proxy. Configures SSL encryption when
    needed.
    """
    def __init__(self, *args, **kwargs):
        """
        Initialize self. Uses the same arguments as
        http.server.BaseHTTPRequestHandler's constructor but also expect a
        `certdir` keyword argument with appropriate path.
        """
        self.certdir = Path(kwargs.pop('certdir')).resolve()
        super().__init__(*args, **kwargs)

    def log_error(self, *args, **kwargs):
        """
        Like log_error in http.server.BaseHTTPRequestHandler but suppresses
        "Request timed out: timeout('timed out',)".
        """
        if not isinstance(args[0], socket.timeout):
            super().log_error(*args, **kwargs)

    def get_cert(self, hostname):
        """
        If needed, generate a signed x509 certificate for `hostname`. Return
        paths to certificate's key file and to certificate itself in a tuple.
        """
        root_keyfile  = self.certdir / 'rootCA.key'
        root_certfile = self.certdir / 'rootCA.pem'
        keyfile       = self.certdir / 'site.key'
        certfile      = self.certdir / f'{hostname}.crt'

        with lock:
            requestfile = self.certdir / f'{hostname}.csr'
            if not certfile.exists():
                subprocess.run([
                    'openssl', 'req', '-new', '-key', str(keyfile),
                    '-subj', f'/CN={hostname}', '-out', str(requestfile)
                ], check=True)
                subprocess.run([
                    'openssl', 'x509', '-req', '-in', str(requestfile),
                    '-CA', str(root_certfile), '-CAkey', str(root_keyfile),
                    '-CAcreateserial', '-out', str(certfile), '-days', '1024'
                ], check=True)

        return keyfile, certfile

    def do_CONNECT(self):
        """Wrap the connection with SSL using on-demand signed certificate."""
        hostname = self.path.split(':')[0]
        sslargs = {'server_side': True}
        sslargs['keyfile'], sslargs['certfile'] = self.get_cert(hostname)

        self.send_response(200)
        self.end_headers()

        self.connection = ssl.wrap_socket(self.connection, **sslargs)
        self.rfile = self.connection.makefile('rb', self.rbufsize)
        self.wfile = self.connection.makefile('wb', self.wbufsize)

        connection_header = self.headers.get('Proxy-Connection', '').lower()
        self.close_connection = int(connection_header == 'close')

    def do_GET(self):
        content_length = int(self.headers.get('Content-Length', 0))
        req_body = self.rfile.read(content_length) if content_length else None

        if self.path[0] == '/':
            secure = 's' if isinstance(self.connection, ssl.SSLSocket) else ''
            self.path = f'http{secure}://{self.headers["Host"]}{self.path}'

        self.handle_request(req_body)

    do_OPTIONS = do_DELETE = do_PUT = do_HEAD = do_POST = do_GET

    def handle_request(self, req_body):
        """Default handler that does nothing. Please override."""
        pass


class ThreadingHTTPServer(ThreadingMixIn, HTTPServer):
    """The actual proxy server"""
    address_family, daemon_threads = socket.AF_INET6, True

    def handle_error(self, request, client_address):
        """
        Like handle_error in http.server.HTTPServer but suppresses socket/ssl
        related errors.
        """
        cls, e = sys.exc_info()[:2]
        if not (cls is socket.error or cls is ssl.SSLError):
            return super().handle_error(request, client_address)