Python txtorcon.TorControlProtocol类代码示例

class FakeEndpointAnswers:
    implements(IStreamClientEndpoint)

    def __init__(self, answers):
        self.answers = answers
        # since we use pop() we need these to be "backwards"
        self.answers.reverse()

    def get_info_raw(self, keys):
        ans = ''
        for k in keys.split():
            if len(self.answers) == 0:
                raise TorProtocolError(551, "ran out of answers")
            ans += '%s=%s\r\n' % (k, self.answers.pop())
        return ans[:-2]                 # don't want trailing \r\n

    def get_info_incremental(self, key, linecb):
        linecb('%s=%s' % (key, self.answers.pop()))
        return defer.succeed('')

    def connect(self, protocol_factory):
        self.proto = TorControlProtocol()
        self.proto.transport = proto_helpers.StringTransport()
        self.proto.get_info_raw = self.get_info_raw
        self.proto.get_info_incremental = self.get_info_incremental
        self.proto._set_valid_events('GUARD STREAM CIRC NS NEWCONSENSUS ORCONN NEWDESC ADDRMAP STATUS_GENERAL')

        return defer.succeed(self.proto)

class ProtocolIntegrationTests(unittest.TestCase):
    """
    Tests which use a real TorControlProtocol objects, not a mock.
    """

    def setUp(self):
        self.protocol = TorControlProtocol(lambda: defer.succeed('foo'))
        self.transport = proto_helpers.StringTransport()

    def send(self, line):
        self.protocol.dataReceived(line.strip() + "\r\n")

    @defer.inlineCallbacks
    def test_with_arg(self):
        info = TorInfo(self.protocol)
        pb = info.post_bootstrap

        ## now we hook up the protocol like it connected to a real Tor
        self.protocol.makeConnection(self.transport)

        ## answer all the requests generated by TorControlProtocol boostrapping etc.
        self.send('250-AUTH METHODS=PASSWORD')
        self.send('250 OK')

        ## response to AUTHENTICATE
        self.send('250 OK')

        ## now we're in _bootstrap() in TorControlProtocol()
        self.send("250-version=foo")
        self.send("250 OK")

        self.send("250-events/names=")
        self.send("250 OK")

        self.send("250 OK")  # for USEFEATURE

        ## do the TorInfo magic
        self.send('250-info/names=')
        self.send('250-multi/path/arg/* a documentation string')
        self.send('250 OK')

        ## we had to save this up above due to the "interesting" way
        ## TorInfo switches to become a possible-nice magic thingy
        ## that does attribute-access for you.
        yield pb

        self.assertTrue(hasattr(info, 'multi'))
        self.assertTrue(hasattr(getattr(info, 'multi'), 'path'))
        self.assertTrue(hasattr(getattr(getattr(info, 'multi'), 'path'), 'arg'))

        ## Finally! The test! We see if we can get this multi-path
        ## value with an argument...
        ## a "real" tor example is "net/listeners/socks" which shows
        ## up in info/names as "net/listeners/*"

        d = info.multi.path.arg('quux')
        d.addCallback(CheckAnswer(self, 'foo'))
        self.send("250-multi/path/arg/quux=foo")
        self.send("250 OK")
        yield d

class FakeTorState(object):

    def __init__(self, routers):
        self.routers = routers
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.protocol.transport = proto_helpers.StringTransport()
        self.protocol.makeConnection(self.protocol.transport)

    def _find_circuit_after_extend(self, x):
        return defer.succeed(None)

    def build_circuit(self, routers=None, using_guards=True):
        cmd = "EXTENDCIRCUIT 0 "
        first = True
        for router in routers:
            if first:
                first = False
            else:
                cmd += ','
            if isinstance(router, basestring) and len(router) == 40 \
               and hashFromHexId(router):
                cmd += router
            else:
                cmd += router.id_hex[1:]
        d = self.protocol.queue_command(cmd)
        print "d %r" % (d,)
        d.addCallback(self._find_circuit_after_extend)
        return d

 def test_set_conf_wrong_args(self):
     ctl = TorControlProtocol()
     d = ctl.set_conf('a')
     self.assertTrue(d.called)
     self.assertTrue(d.result)
     self.assertTrue('even number' in d.result.getErrorMessage())
     ## ignore the error so trial doesn't get unhappy
     d.addErrback(lambda foo: True)
     return d

class LogicTests(unittest.TestCase):
    def setUp(self):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.protocol.makeConnection(self.transport)

    def test_set_conf_wrong_args(self):
        ctl = TorControlProtocol()
        d = ctl.set_conf("a")
        self.assertTrue(d.called)
        self.assertTrue(d.result)
        self.assertTrue("even number" in d.result.getErrorMessage())
        ## ignore the error so trial doesn't get unhappy
        d.addErrback(lambda foo: True)
        return d

 def setUp(self):
     self.protocol = TorControlProtocol()
     self.protocol.connectionMade = lambda: None
     self.transport = proto_helpers.StringTransportWithDisconnection()
     self.protocol.makeConnection(self.transport)
     # why doesn't makeConnection do this?
     self.transport.protocol = self.protocol

    def connect(self, protocol_factory):
        self.proto = TorControlProtocol()
        self.proto.transport = proto_helpers.StringTransport()
        self.proto.get_info_raw = self.get_info_raw
        self.proto.get_info_incremental = self.get_info_incremental
        self.proto._set_valid_events('GUARD STREAM CIRC NS NEWCONSENSUS ORCONN NEWDESC ADDRMAP STATUS_GENERAL')

        return defer.succeed(self.proto)

class FakeReactorTcp(FakeReactor):
    implements(IReactorTCP)

    failures = 0
    _port_generator = port_generator()

    def __init__(self, test):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.transport = FakeProcessTransport()
        self.transport.protocol = self.protocol

        def blam():
            self.protocol.outReceived("Bootstrap")
        self.transport.closeStdin = blam
        self.protocol.makeConnection(self.transport)
        FakeReactor.__init__(self, test, self.transport, lambda x: None)

    def listenTCP(self, port, factory, **kwargs):
        '''returns IListeningPort'''
        if self.failures > 0:
            self.failures -= 1
            raise error.CannotListenError(None, None, None)

        if port == 0:
            port = self._port_generator.next()
        p = FakeListeningPort(port)
        p.factory = factory
        p.startListening()
        return p

    def connectTCP(self, host, port, factory, timeout, bindAddress):
        '''should return IConnector'''
#        print "CONN", host, port, factory
#        print dir(factory)
#        print "XX", factory
        r = tcp.Connector(host, port, factory, timeout, bindAddress, reactor=self)
#        print dir(r)

        def blam(*args):
            print "BLAAAAAM", args
        r.connect = blam
        return r

    def __init__(self, test):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.transport.protocol = self.protocol

        def blam():
            self.protocol.outReceived(b"Bootstrap")
        self.transport.closeStdin = blam
        self.protocol.makeConnection(self.transport)
        self.test = test

    def __init__(self, test):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.transport = FakeProcessTransport()
        self.transport.protocol = self.protocol

        def blam():
            self.protocol.outReceived("Bootstrap")
        self.transport.closeStdin = blam
        self.protocol.makeConnection(self.transport)
        FakeReactor.__init__(self, test, self.transport, lambda x: None)

class AuthenticationTests(unittest.TestCase):
    def setUp(self):
        self.protocol = TorControlProtocol()
        self.transport = proto_helpers.StringTransport()

    def send(self, line):
        self.protocol.dataReceived(line.strip() + "\r\n")
        
    def test_authenticate_cookie(self):
        self.protocol.makeConnection(self.transport)
        self.assertEqual(self.transport.value(), 'PROTOCOLINFO 1\r\n')
        self.transport.clear()
        cookie_data = 'cookiedata!cookiedata!cookiedata'
        open('authcookie', 'w').write(cookie_data)
        self.send('250-PROTOCOLINFO 1')
        self.send('250-AUTH METHODS=COOKIE,HASHEDPASSWORD COOKIEFILE="authcookie"')
        self.send('250-VERSION Tor="0.2.2.34"')
        self.send('250 OK')

        self.assertEqual(self.transport.value(), 'AUTHENTICATE %s\r\n' % cookie_data.encode("hex"))

    def test_authenticate_password(self):
        self.protocol.password = 'foo'
        self.protocol.makeConnection(self.transport)
        self.assertEqual(self.transport.value(), 'PROTOCOLINFO 1\r\n')
        self.transport.clear()
        self.send('250-PROTOCOLINFO 1')
        self.send('250-AUTH METHODS=HASHEDPASSWORD')
        self.send('250-VERSION Tor="0.2.2.34"')
        self.send('250 OK')

        self.assertEqual(self.transport.value(), 'AUTHENTICATE %s\r\n' % "foo".encode("hex"))

    def confirmAuthFailed(self, *args):
        self.auth_failed = True
        
    def test_authenticate_no_password(self):
        self.protocol.post_bootstrap.addErrback(self.confirmAuthFailed)
        self.auth_failed = False
        
        self.protocol.makeConnection(self.transport)
        self.assertEqual(self.transport.value(), 'PROTOCOLINFO 1\r\n')
        
        self.send('250-PROTOCOLINFO 1')
        self.send('250-AUTH METHODS=HASHEDPASSWORD')
        self.send('250-VERSION Tor="0.2.2.34"')
        self.send('250 OK')

        self.assertTrue(self.auth_failed)

class TorStatePy3Tests(unittest.TestCase):

    def setUp(self):
        self.protocol = TorControlProtocol()
        self.state = TorState(self.protocol)
        # avoid spew in trial logs; state prints this by default
        self.state._attacher_error = lambda f: f
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.protocol.makeConnection(self.transport)

    def send(self, line):
        self.protocol.dataReceived(line.strip() + b"\r\n")

    def test_attacher_coroutine(self):
        @implementer(IStreamAttacher)
        class MyAttacher(object):

            def __init__(self, answer):
                self.streams = []
                self.answer = answer

            async def attach_stream(self, stream, circuits):
                self.streams.append(stream)
                x = await defer.succeed(self.answer)
                return x

        self.state.circuits[1] = FakeCircuit(1)
        self.state.circuits[1].state = 'BUILT'
        attacher = MyAttacher(self.state.circuits[1])
        self.state.set_attacher(attacher, FakeReactor(self))

        # boilerplate to finish enough set-up in the protocol so it
        # works
        events = 'GUARD STREAM CIRC NS NEWCONSENSUS ORCONN NEWDESC ADDRMAP STATUS_GENERAL'
        self.protocol._set_valid_events(events)
        self.state._add_events()
        for ignored in self.state.event_map.items():
            self.send(b"250 OK")

        self.send(b"650 STREAM 1 NEW 0 ca.yahoo.com:80 SOURCE_ADDR=127.0.0.1:54327 PURPOSE=USER")
        self.send(b"650 STREAM 1 REMAP 0 87.248.112.181:80 SOURCE=CACHE")
        self.assertEqual(len(attacher.streams), 1)
        self.assertEqual(attacher.streams[0].id, 1)
        self.assertEqual(len(self.protocol.commands), 1)
        self.assertEqual(self.protocol.commands[0][1], b'ATTACHSTREAM 1 1')

class ProtocolTests(unittest.TestCase):
    def setUp(self):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransport()
        self.protocol.makeConnection(self.transport)

    def tearDown(self):
        self.protocol = None

    def send(self, line):
        self.protocol.dataReceived(line.strip() + "\r\n")

    def test_statemachine_broadcast_no_code(self):
        try:
            self.protocol._broadcast_response("foo")
            self.fail()
        except RuntimeError, e:
            self.assertTrue("No code set yet" in str(e))

class DisconnectionTests(unittest.TestCase):
    def setUp(self):
        self.protocol = TorControlProtocol()
        self.protocol.connectionMade = lambda: None
        self.transport = proto_helpers.StringTransportWithDisconnection()
        self.protocol.makeConnection(self.transport)
        # why doesn't makeConnection do this?
        self.transport.protocol = self.protocol

    def tearDown(self):
        self.protocol = None

    def test_disconnect_callback(self):
        """
        see that we get our callback on_disconnect if the transport
        goes away
        """
        def it_was_called(*args):
            it_was_called.yes = True
            return None
        it_was_called.yes = False
        self.protocol.on_disconnect.addCallback(it_was_called)
        self.protocol.on_disconnect.addErrback(it_was_called)
        f = failure.Failure(error.ConnectionDone("It's all over"))
        self.protocol.connectionLost(f)
        self.assertTrue(it_was_called.yes)

    def test_disconnect_errback(self):
        """
        see that we get our callback on_disconnect if the transport
        goes away
        """
        def it_was_called(*args):
            it_was_called.yes = True
            return None
        it_was_called.yes = False
        self.protocol.on_disconnect.addCallback(it_was_called)
        self.protocol.on_disconnect.addErrback(it_was_called)
        f = failure.Failure(RuntimeError("The thing didn't do the stuff."))
        self.protocol.connectionLost(f)
        self.assertTrue(it_was_called.yes)

 def setUp(self):
     self.protocol = TorControlProtocol()
     self.transport = proto_helpers.StringTransport()