def encode(*args):
    """Encode given values using AMF0.

    @rtype: VecBuf
    """
    vb = VecBuf()
    _encode(vb, args)
    return vb def decode(data):
    """Decode AMF0-encoded buffer of data.

    @type data: VecBuf

    @returns: list of objects
    """
    try:
        return _decode(data)
    except VecBufEOB:
        raise DecoderError('Incomplete encoded data') class Object(OrderedObject):
    """AMF Object class, with chronological attribute ordering, for
    increased compatibility.
    """ class Header(object):
class ChunkStreamParseError(ValueError):
class ChunkStreamValueError(ValueError):
class Demuxer(object):
class Chunker(object):
class SimpleChunkProducer(object):
class Muxer(object):
    def __init__(self, cs_id, time, size, type, ms_id, **kw):
    def __repr__(self):
        def _repr_maybe(v, provided):
def absolutize(h, base):
def read_header_head(s):
    def __init__(self, protocol):
    def doSetChunkSize(self, header, new_size):
    def doAbortMessage(self, header, cs_id):
    def doACK(self, header, seq_num):
    def doUserControlMessage(self, header, evt_type, body):
    def doWindowSize(self, header, window_size):
    def doSetBandwidth(self, header, window_size, limit_type):
    def controlMessageUnknown(self, header, body):
    def controlMessageReceived(self, header, body):
    def gen_handler(self):
def _encode_basic_header(h_type, cs_id, time=None):
def encode_full_header(cs_id, time, size, msg_type, ms_id):
def encode_comp_header(h_type, cs_id, time, size, msg_type, ms_id):
    def set_chunk_size(self, new_chunk_size):
    def __call__(self, cs_id, first_header, body, time=None):
    def __init__(self, transport):
    def queue_chunker(self, priority, chunker):
    def sync(self, priority):
    def pauseProducing(self):
    def resumeProducing(self):
    def stopProducing(self):
    def __init__(self, transport):
    def _precompute_dispatch(self, amf_ver):
    def _make_adhoc_csid(self, mt):
    def set_chunk_size(self, new_chunk_size):
    def sendMessage(self, time, type_, ms_id, body, absolute=False):
DEFAULT_CHUNK_SIZE = 128
(PROTO_SET_CHUNK_SIZE, PROTO_ABORT_MESSAGE, PROTO_ACK, PROTO_USER_CONTROL,
 PROTO_WINDOW_SIZE, PROTO_SET_BANDWIDTH, MSG_COMMAND, MSG_DATA, MSG_SO,
 MSG_AUDIO, MSG_VIDEO, MSG_AGGREGATE) = range(12) RTMP_CHUNK_SIZE = 0x01
RTMP_ABORT = 0x02
RTMP_ACK = 0x03
RTMP_USER_CONTROL = 0x04
RTMP_WINDOW_SIZE = 0x05
RTMP_SET_BANDWIDTH = 0x06
RTMP_AUDIO = 0x08
RTMP_VIDEO = 0x09
RTMP_DATA = 0x12
RTMP_SHARED_OBJ = 0x13
RTMP_COMMAND = 0x14
RTMP_AGGREGATE = 0x16
RTMP_DATA_AMF3 = 0x0f
RTMP_SHARED_OBJ_AMF3 = 0x10
RTMP_COMMAND_AMF3 = 0x11
UCTRL_STREAM_BEGIN = 0x0
UCTRL_STREAM_EOF = 0x1
UCTRL_STREAM_DRY = 0x2
UCTRL_BUFFER_LENGTH = 0x3
UCTRL_STREAM_RECORDED = 0x4
UCTRL_PING = 0x6
UCTRL_PONG = 0x7
AF_PCM_PLATFORM = 0
AF_ADPCM = 1
AF_MP3 = 2
AF_PCM_LITTLE_ENDIAN = 3
AF_NELLYMOSER_16K = 4
AF_NELLYMOSER_8K = 5
AF_NELLYMOSER = 6
AF_ALAW = 7
AF_ULAW = 8
AF_AAC = 10
AF_SPEEX = 11
AF_MP3_8K = 14
AF_DEVICE_SPECIFIC = 15
VF_JPEG = 1
VF_H263 = 2
VF_SCREEN = 3
VF_VP6 = 4
VF_VP6_ALPHA = 5
VF_SCREEN_2 = 6
VF_H264 = 7
FT_VIDEO_KEYFRAME = 1
FT_VIDEO_INTERFAME = 2
FT_VIDEO_DISPOSABLE_INTERFRAME = 3
FT_VIDEO_GENERATED_KEYFRAME = 4
FT_VIDEO_INFO = 5
FT_AAC_HEADER = 0
FT_AAC_DATA = 1
FT_H264_HEADER = 0
FT_H264_NALU = 1
FT_H264_EOS = 2 class CommandDispatchProtocol(DispatchProtocol):

    def __init__(self):
        DispatchProtocol.__init__(self)

        self._cc_queue = CancellableCallQueue()
        self._call_tracker = DeferredTracker()

    def doCommand(self, ts, ms_id, args):
        cmd = args[0]

        handler_m = getattr(self, 'command_%s' % (cmd,), None)

        if handler_m is None:
            self._cc_queue.callLater(0, self.unknownCommandType, cmd, ts,
                                     ms_id, args[1:])
        else:
            self._cc_queue.callLater(0, self._handler_wrapper, handler_m,
                                     ts, ms_id, args[1:])

    def _handler_wrapper(self, handler, ts, ms_id, args):
        # wrap in try/except...?
        handler(ts, ms_id, *args)

    def command__result(self, ts, ms_id, trans_id, *args):
        d = self._call_tracker.pop_deferred(ms_id, trans_id)

        if d:
            d.callback(args)
        else:
            self.unexpectedCallResult(ts, ms_id, trans_id, args)

    def command__error(self, ts, ms_id, trans_id, *args):
        d = self._call_tracker.pop_deferred(ms_id, trans_id)

        if d:
            d.errback(failure.Failure(CommandResultError(*args)))
        else:
            self.unexpectedCallError(ts, ms_id, trans_id, args)

    def unknownCommandType(self, cmd, ts, msid, args):
        raise NotImplementedError('unknown command %r%r' % (cmd,
                                                            (ts, msid, args)))

    def unexpectedCallResult(self, ts, ms_id, trans_id, args):
        log.warning('unexpected _result: at %r, stream %r, trans %r, args: %r',
                    ts, ms_id, trans_id, args)

    def unexpectedCallError(self, ts, ms_id, trans_id, args):
        log.warning('unexpected _error: at %r, stream %r, trans %r, args: %r',
                    ts, ms_id, trans_id, args)


    def connectionLost(self, reason=protocol.connectionDone):
        self._cc_queue.cancel_all()
        pending_calls = list(self._call_tracker.iter_all())
        self._call_tracker.clear()
        for ms_id, d in pending_calls:
            d.errback(reason)
        DispatchProtocol.connectionLost(self, reason)

    def encode_amf(self, *args):
        # for now only supporting AMF0
        return amf0.encode(*args)

    def _send_command(self, ts, ms_id, body, track_id):
        # track_id > 0 -> will return a tracking deferred
        d = None

        if track_id:
            d = defer.Deferred()
            self._call_tracker.push_deferred(ms_id, track_id, d)

        self.muxer.sendMessage(0, chunks.MSG_COMMAND, ms_id, body)
        return d

    def _sendRemote(self, ms_id, cmd, args, kwargs, track):
        # ignoring kwargs for now...
        trans_id = 0
        if track:
            trans_id = self._call_tracker.next_trans(ms_id)
        encoded_args = self.encode_amf(cmd, trans_id, *args)

        # hardcoding 0 time, does not seem to matter much...
        return self._send_command(0, ms_id, encoded_args, trans_id)

    def callRemote(self, ms_id, cmd, *args, **kw):
        return self._sendRemote(ms_id, cmd, args, kw, True)

    def signalRemote(self, ms_id, cmd, *args, **kw):
        # similar to callRemote, except we don't expect any results
        return self._sendRemote(ms_id, cmd, args, kw, False) class EventDispatchProtocol(CommandDispatchProtocol):
    def __init__(self):
        CommandDispatchProtocol.__init__(self)

        self._events = StatusEventTracker()

    def _onStatus_ev_key(self, ms_id):
        return (None, ms_id)

    def waitStatus(self, ms_id, code):
        return self._events.wait(self._onStatus_ev_key(ms_id), code)

    def command_onStatus(self, ts, ms_id, _trans_id, _none, info):
        # trans_id not used, and _none seems to always be None...

        def miss_handler():
            self.unhandledOnStatus(ts, ms_id, info)

        self._events.dispatch(self._onStatus_ev_key(ms_id), info,
                              miss_h=miss_handler)

    def unhandledOnStatus(self, ts, ms_id, info):
        log.warning('unhandled onStatus: at %r, stream %r, info: %r',
                    ts, ms_id, info)

    def connectionLost(self, reason=protocol.connectionDone):
        self._events.cancel_all(reason=reason)

        CommandDispatchProtocol.connectionLost(self, reason) class CallDispatchProtocol(EventDispatchProtocol):
    def __init__(self):
        EventDispatchProtocol.__init__(self)

    def session_time(self):
        return time.time() - self.session_init_time

    def unknownCommandType(self, cmd, ts, ms_id, args):
        trans_id = args[0]

        handler_m = getattr(self, 'remote_%s' % (cmd,), None)

        if handler_m is None:
            d = defer.maybeDeferred(self.unknownRemoteCall, cmd, ts, ms_id,
                                    args[1:])
        else:
            d = defer.maybeDeferred(handler_m, ts, ms_id, *args[1:])

        if trans_id:
            d.addCallback(self._remote_handler_cb, ms_id, trans_id)

        d.addErrback(self._remote_abort_handler_eb)
        d.addErrback(self._remote_handler_eb, ms_id, trans_id)

    def _remote_abort_handler_eb(self, failure):
        failure.trap(CallAbortedException)
        # log failure but do nothing more
        log.debug('remote call aborted: %s', failure.value)

    def _remote_handler_cb(self, result, ms_id, trans_id):
        # log.debug('remote call result: %r', result)
        if not isinstance(result, (tuple, list)):
            result = (result,)

        body = self.encode_amf('_result', trans_id, *result)

        ts = ms_time_wrapped(self.session_time())
        self.muxer.sendMessage(ts, chunks.MSG_COMMAND, ms_id, body)

    def _remote_handler_eb(self, failure, ms_id, trans_id):
        if log.isEnabledFor(logging.DEBUG):
            log.info('remote call failure: %s', failure.value,
                     exc_info=(failure.type, failure.value,
                               failure.getTracebackObject()))
        else:
            log.info('remote call failure: %s', failure.value)

        fatal = False
        if failure.check(CallResultError):
            body = self.encode_amf('_error', trans_id,
                                   *failure.value.get_error_args())
            fatal = failure.value.is_fatal
        else:
            err = amf0.Object(code='NetStream.Failed', level='error',
                              description=repr(failure.value))
            body = self.encode_amf('_error', trans_id, None, err)

        ts = ms_time_wrapped(self.session_time())
        self.muxer.sendMessage(ts, chunks.MSG_COMMAND, ms_id, body)

        if fatal:
            self.transport.loseConnection()

    def unknownRemoteCall(self, cmd, ts, ms_id, args):
        # seems that we're just supposed to silently ignore the request
        log.warning('unknown method called: %s, args: %r', cmd, args)
        raise CallAbortedException('unknown command %r' % (cmd,)) class UnexpectedStatusError(ValueError):
    pass class ProtocolContractError(ValueError):
    pass class CommandResultError(RuntimeError):
    def __str__(self):
        desc = None
        if (self.args and len(self.args) == 2 and
            isinstance(self.args[1], amf0.Object)):
            obj = self.args[1]
            if obj.has_key('description'):
                desc = '%s: %s' % (obj.get('code', self.__class__.__name__),
                                   obj.description)
        return desc or RuntimeError.__str__(self) class CallAbortedException(Exception):
    """Call aborted""" def vb(s):
    # TODO: add this functionality to the vecbuf/VecBuf itself
    return VecBuf([s]) def vb_clone(vb):
    # TODO: add this functionality to the vecbuf/VecBuf itself
    return VecBuf(vb.peek_seq(len(vb))) class _ProtocolTestBase(unittest.TestCase):
    def setUp(self):
        self.messages = []

    def clear(self):
        self.messages[:] = []


    def make_protocol_class(self):
        raise NotImplementedError()

    def make_factory(self):
        return TestFactory()

    def make_protocol(self):
        f = self.make_factory()
        f.protocol = self.make_protocol_class()

        p = f.buildProtocol(None)

        return p

    def connect_protocol(self, protocol):
        t = StringTransport()
        protocol.makeConnection(t)

        # ... and let's skip the handshake
        protocol.dataReceived('.')

        return t

    def build_proto(self):
        p = self.make_protocol()
        t = self.connect_protocol(p)
        return p, t, p._demuxer, p.muxer def mixin(base_class, mixin_class, name=None):
    if not name:
        name = 'Mixed' + base_class.__name__
    return type(name, (mixin_class, base_class), {}) class TestDemuxerMixin(object):
    """Allow injecting complete ready messages as if they were demuxed
    by the demuxer.
    """
    def inject(self, cs_id, time, type_, ms_id, body):
        h = chunks.Header(cs_id, time, len(body), type_, ms_id)
        if cs_id == 2 and 0 < type_ < 8 and ms_id == 0:
            self.controlMessageReceived(h, body)
        else:
            self.protocol.messageReceived(h, body)

    def sinject(self, cs_id, time, type_, ms_id, data):
        body = vb(data)
        self.inject(cs_id, time, type_, ms_id, body) class TestMuxer(object):
    def __init__(self, transport):
        self.transport = transport
        self.chunk_size = chunks.DEFAULT_CHUNK_SIZE
        self.messages = []

    def set_chunk_size(self, new_chunk_size):
        self.chunk_size = new_chunk_size

    def sendMessage(self, time, type_, ms_id, body, absolute=False):
        self.messages.append((time, type_, ms_id, body.read(len(body)),
                              absolute)) class TestHandshaker(object):
    def __init__(self, protocol, epoch_base, is_client=False):
        self.protocol = protocol

    def gen_handler(self):
        s = yield 1
        s.read(1)
        self.protocol.handshakeSucceeded(0, 0) class ArtificialError(RuntimeError):
    """Purely artificial test error""" class ArtificialRemoteError(CallResultError):
    """Completely artificial test error"""
    code = 'Artificially.Failed' def muxer_messages(mux):
    return [(m[0], m[1], m[2],
             (decode_amf(vb(m[3]))
              if m[1] in (chunks.MSG_COMMAND, chunks.MSG_DATA)
              else m[3]),
             m[4])
            for m in mux.messages] def unvb(vb):
    return vb.read(len(vb))[:]
from twisted.internet import defer, error, reactor
from twimp.amf0 import encode as encode_amf, decode as decode_amf, Object
from twimp import chunks
from twimp import const
from twimp.dispatch import CommandDispatchProtocol
from twimp.dispatch import EventDispatchProtocol
from twimp.dispatch import CallDispatchProtocol
from twimp.error import UnexpectedStatusError, ProtocolContractError
from twimp.error import CommandResultError
from twimp.error import CallAbortedException
from twimp.helpers import vb, vb_clone
from test.test_proto import _ProtocolTestBase, mixin
from test.test_proto import TestDemuxerMixin, TestMuxer, TestHandshaker
from test.common import ArtificialError, ArtificialRemoteError
from test.helpers import muxer_messages, unvb
#   Copyright (c) 2011  Arek Korbik
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.








class TestCommandDispatchProtocol(_ProtocolTestBase):
    timeout = 10

    # note: CDP tests ideally should be performed with a controllable
    # reactor, but the CDP code would need to be adapted, so for now
    # just relying on ordered scheduling of callLater() calls...

    def make_protocol_class(self):
        tself = self
        CDP = CommandDispatchProtocol
        class TCommandDispatchProtocol(CDP):
            demuxer_class = mixin(CDP.demuxer_class, TestDemuxerMixin)
            muxer_class = TestMuxer

            handshaker_class = TestHandshaker