tor_control_socket_test.py

import os
import sys
import socket
import tempfile
import unittest
import threading

from cffi import FFI

LIB_PATH = "libtor_api.tor_api_owning_control.dylib"

def load_lib():
    ffi = FFI()
    # This is the output of: `cc -E tor_api.h`
    ffi.cdef("""typedef struct tor_main_configuration_t tor_main_configuration_t;
typedef int tor_control_socket_t;
tor_main_configuration_t *tor_main_configuration_new(void);
# 49 "tor_api.h"
int tor_main_configuration_set_command_line(tor_main_configuration_t *cfg,
                                            int argc, char *argv[]);
void tor_main_configuration_free(tor_main_configuration_t *cfg);
# 86 "tor_api.h"
int tor_run_main(const tor_main_configuration_t *);
const char *tor_api_get_provider_version(void);
tor_control_socket_t tor_main_configuration_setup_control_socket(
                                           tor_main_configuration_t *cfg);
# 98 "tor_api.h"
int tor_main(int argc, char **argv);""")
    lib = ffi.dlopen(LIB_PATH)
    return ffi, lib

ffi, lib = load_lib()

class TestTorAPI(unittest.TestCase):
    def test_ctrl_getinfo(self):
        """
        This checks we can run the GETINFO command via the ctrl socket
        """
        cfg = lib.tor_main_configuration_new()
        ctrl_fd = lib.tor_main_configuration_setup_control_socket(cfg)
        ctrl_socket = socket.socket(fileno=ctrl_fd)

        main_thread = threading.Thread(target=lib.tor_run_main, args=(cfg, ))
        main_thread.start()

        ctrl_socket.send(b"GETINFO version\r\n")
        in_data = ctrl_socket.recv(1024)

        self.assertTrue(
            in_data.startswith(b"250-version=")
        )

        ctrl_socket.close()
        main_thread.join()
        lib.tor_main_configuration_free(cfg)

    def test_ctrl_close_exits(self):
        """
        This checks that by closing the fd tor will exit cleanly
        """
        cfg = lib.tor_main_configuration_new()
        ctrl_fd = lib.tor_main_configuration_setup_control_socket(cfg)
        ctrl_socket = socket.socket(fileno=ctrl_fd)

        # XXX we should maybe check that the return code of this is good
        main_thread = threading.Thread(target=lib.tor_run_main, args=(cfg, ))
        main_thread.start()

        ctrl_socket.close()

        main_thread.join()
        lib.tor_main_configuration_free(cfg)

if __name__ == "__main__":
    unittest.main()