blender-benchmark-bundle/benchmark/submission/auth.py

import http.server
import logging
import random
import urllib.parse

from . import html, sockutil

log = logging.getLogger(__name__)


class TokenHTTPHandler(http.server.BaseHTTPRequestHandler):
    """Handle GET requests with tokens on the URL."""

    def do_GET(self):
        # /?token=72157630789362986-5405f8542b549e95

        qs = urllib.parse.urlsplit(self.path).query
        url_vars = urllib.parse.parse_qs(qs)

        self.server.auth_token = url_vars['token'][0]
        assert (isinstance(self.server.auth_token, str))

        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()

        self.wfile.write(html.auth_okay_html)


class TokenHTTPServer(http.server.HTTPServer):
    """HTTP server on a random port, which will receive the token."""

    def __init__(self) -> None:
        self.log = log.getChild('TokenHTTPServer')

        self.local_addr = self.listen_port()
        self.log.info('Creating HTTP server at %s', self.local_addr)
        self.auth_token = None

        http.server.HTTPServer.__init__(self, self.local_addr, TokenHTTPHandler)

    def listen_port(self):
        """Returns the hostname and TCP/IP port number to listen on.

        Finds a random free port between 1100 and 20000.
        """

        # Find a random free port
        local_addr = ('localhost', int(random.uniform(1100, 20000)))
        self.log.debug('Finding free port starting at %s', local_addr)
        return sockutil.find_free_port(local_addr)

    def wait_for_token(self, timeout: float):
        """Starts the HTTP server, waits for the Token."""

        if self.auth_token is None:
            self.timeout = timeout
            self.handle_request()

        if self.auth_token:
            self.log.info('Auth token received: %s' % self.auth_token)

        return self.auth_token

    @property
    def auth_callback_url(self) -> str:
        return f'http://localhost:{self.local_addr[1]}/'