v3_onion_test.py

#!/usr/bin/env python3.6
# -*- coding: utf-8 -*-


import base64
import binascii
import hashlib
import os
import random
import shlex
import socket
import subprocess
import tempfile
import time
import typing

import stem
from stem.control import Controller

if typing.TYPE_CHECKING:
    from tempfile import TemporaryDirectory


def get_available_port(min_port: int, max_port: int) -> str:
    with socket.socket() as tmpsock:
        while True:
            try:
                tmpsock.bind(('127.0.0.1', random.randint(min_port, max_port)))
                break
            except OSError:
                pass
        _, port = tmpsock.getsockname()
    return port


class Tor(object):

    def __init__(self) -> None:
        self.tor_process        = None  # type: subprocess.Popen
        self.controller         = None  # type: Controller
        self.tor_data_directory = None  # type: TemporaryDirectory

    def connect(self, port: str) -> bool:
        self.tor_data_directory = tempfile.TemporaryDirectory()

        torrc_data = """\
       DataDirectory {{data_directory}}
       SocksPort {{socks_port}}
       ControlSocket {{control_socket}}
       CookieAuthentication 1
       CookieAuthFile {{cookie_auth_file}}
       AvoidDiskWrites 1
       Log notice stdout
       GeoIPFile /usr/share/tor/geoip
       GeoIPv6File /usr/share/tor/geoip6
       """

        tor_control_socket   = os.path.join(self.tor_data_directory.name, 'control_socket')
        tor_cookie_auth_file = os.path.join(self.tor_data_directory.name, 'cookie')
        tor_torrc            = os.path.join(self.tor_data_directory.name, 'torrc')

        torrc_data = torrc_data.replace('{{data_directory}}', self.tor_data_directory.name)
        torrc_data = torrc_data.replace('{{socks_port}}', str(port))
        torrc_data = torrc_data.replace('{{control_socket}}', str(tor_control_socket))
        torrc_data = torrc_data.replace('{{cookie_auth_file}}', tor_cookie_auth_file)

        with open(tor_torrc, 'w') as f:
            f.write(torrc_data)

        start_ts = time.monotonic()

        self.tor_process = subprocess.Popen(['{}/stemtest/tor/src/or/tor'.format(os.getenv("HOME")), '-f', tor_torrc],
                                            stdout=subprocess.PIPE,
                                            stderr=subprocess.PIPE)
        time.sleep(2)
        self.controller = Controller.from_socket_file(path=tor_control_socket)
        self.controller.authenticate()

        while True:
            time.sleep(0.1)
            response  = self.controller.get_info("status/bootstrap-phase")
            res_parts = shlex.split(response)
            progress  = res_parts[2].split('=')[1]
            summary   = res_parts[4].split('=')[1]
            print("{}: {}% - {}{}".format('connecting_to_tor', progress, summary, "\033[K"), end="\r")

            if summary == 'Done':
                print()
                break
            if time.monotonic() - start_ts > 45:
                return False
        return True

    def stop(self) -> None:
        if self.tor_process:
            self.tor_process.terminate()
            time.sleep(0.2)
            if not self.tor_process.poll():
                self.tor_process.kill()


def main() -> None:
    try:
        pids = subprocess.check_output("ps aux |grep '[t]orrc' | awk '{print $2}' 2>/dev/null",
                                       shell=True).split(b'\n')
        for pid in pids:
            subprocess.Popen("kill {}".format(int(pid)), shell=True).wait()
    except ValueError:
        pass

    tor_port = get_available_port(1000, 65535)
    tor = Tor()
    if not tor.connect(tor_port):
        print("\nTor timed out!")
        exit(1)

    print('Starting Onion Service... ', end='', flush=True)
    try:
        service = tor.controller.create_ephemeral_hidden_service(ports={80: 5000},
                                                                 key_type="NEW",
                                                                 key_content="ED25519-V3",
                                                                 await_publication=True)
    except stem.OperationFailed as e:
        print(f"OperationFailed:\n{e}")
        tor.stop()
        exit(1)

    print('Onion Service {} is now online'.format(service.service_id))
    tor.controller.remove_hidden_service(service.service_id)
    tor.stop()
    exit(0)


if __name__ == '__main__':
    main()