Permalink
Switch branches/tags
Nothing to show
Find file
765480b Aug 25, 2012
548 lines (475 sloc) 19 KB
"""
Provides classes for creating RTMP (Real Time Message Protocol) servers and
clients.
"""
import pyamf.amf0
import pyamf.util.pure
import rtmp_protocol_base
import socket
import logging
class FileDataTypeMixIn(pyamf.util.pure.DataTypeMixIn):
"""
Provides a wrapper for a file object that enables reading and writing of raw
data types for the file.
"""
def __init__(self, fileobject):
self.fileobject = fileobject
pyamf.util.pure.DataTypeMixIn.__init__(self)
def read(self, length):
return self.fileobject.read(length)
def write(self, data):
self.fileobject.write(data)
def flush(self):
self.fileobject.flush()
def at_eof(self):
return False
class DataTypes:
""" Represents an enumeration of the RTMP message datatypes. """
NONE = -1
SET_CHUNK_SIZE = 1
USER_CONTROL = 4
WINDOW_ACK_SIZE = 5
SET_PEER_BANDWIDTH = 6
SHARED_OBJECT = 19
COMMAND = 20
class SOEventTypes:
""" Represents an enumeration of the shared object event types. """
USE = 1
RELEASE = 2
CHANGE = 4
MESSAGE = 6
CLEAR = 8
DELETE = 9
USE_SUCCESS = 11
class UserControlTypes:
""" Represents an enumeration of the user control event types. """
STREAM_BEGIN = 0
STREAM_EOF = 1
STREAM_DRY = 2
SET_BUFFER_LENGTH = 3
STREAM_IS_RECORDED = 4
PING_REQUEST = 6
PING_RESPONSE = 7
class RtmpReader:
""" This class reads RTMP messages from a stream. """
chunk_size = 128
def __init__(self, stream):
"""
Initialize the RTMP reader and set it to read from the specified stream.
"""
self.stream = stream
def __iter__(self):
return self
def next(self):
""" Read one RTMP message from the stream and return it. """
if self.stream.at_eof():
raise StopIteration
# Read the message into body_stream. The message may span a number of
# chunks (each one with its own header).
message_body = []
msg_body_len = 0
header = rtmp_protocol_base.header_decode(self.stream)
# FIXME: this should be really implemented inside header_decode
if header.datatype == DataTypes.NONE:
header = self.prv_header
self.prv_header = header
while True:
read_bytes = min(header.bodyLength - msg_body_len, self.chunk_size)
message_body.append(self.stream.read(read_bytes))
msg_body_len += read_bytes
if msg_body_len >= header.bodyLength:
break
next_header = rtmp_protocol_base.header_decode(self.stream)
# WORKAROUND: even though the RTMP specification states that the
# extended timestamp field DOES NOT follow type 3 chunks, it seems
# that Flash player 10.1.85.3 and Flash Media Server 3.0.2.217 send
# and expect this field here.
if header.timestamp >= 0x00ffffff:
self.stream.read_ulong()
assert next_header.streamId == -1, (header, next_header)
assert next_header.datatype == -1, (header, next_header)
assert next_header.timestamp == -1, (header, next_header)
assert next_header.bodyLength == -1, (header, next_header)
assert header.bodyLength == msg_body_len, (header, msg_body_len)
body_stream = pyamf.util.BufferedByteStream(''.join(message_body))
# Decode the message based on the datatype present in the header
ret = {'msg':header.datatype}
if ret['msg'] == DataTypes.USER_CONTROL:
ret['event_type'] = body_stream.read_ushort()
ret['event_data'] = body_stream.read()
elif ret['msg'] == DataTypes.WINDOW_ACK_SIZE:
ret['window_ack_size'] = body_stream.read_ulong()
elif ret['msg'] == DataTypes.SET_PEER_BANDWIDTH:
ret['window_ack_size'] = body_stream.read_ulong()
ret['limit_type'] = body_stream.read_uchar()
elif ret['msg'] == DataTypes.SHARED_OBJECT:
decoder = pyamf.amf0.Decoder(body_stream)
obj_name = decoder.readString()
curr_version = body_stream.read_ulong()
flags = body_stream.read(8)
# A shared object message may contain a number of events.
events = []
while not body_stream.at_eof():
event = self.read_shared_object_event(body_stream, decoder)
events.append(event)
ret['obj_name'] = obj_name
ret['curr_version'] = curr_version
ret['flags'] = flags
ret['events'] = events
elif ret['msg'] == DataTypes.COMMAND:
decoder = pyamf.amf0.Decoder(body_stream)
commands = []
while not body_stream.at_eof():
commands.append(decoder.readElement())
ret['command'] = commands
#elif ret['msg'] == DataTypes.NONE:
# print 'WARNING: message with no datatype received.', header
# return self.next()
elif ret['msg'] == DataTypes.SET_CHUNK_SIZE:
ret['chunk_size'] = body_stream.read_ulong()
else:
assert False, header
logging.debug('recv %r', ret)
return ret
def read_shared_object_event(self, body_stream, decoder):
"""
Helper method that reads one shared object event found inside a shared
object RTMP message.
"""
so_body_type = body_stream.read_uchar()
so_body_size = body_stream.read_ulong()
event = {'type':so_body_type}
if event['type'] == SOEventTypes.USE:
assert so_body_size == 0, so_body_size
event['data'] = ''
elif event['type'] == SOEventTypes.RELEASE:
assert so_body_size == 0, so_body_size
event['data'] = ''
elif event['type'] == SOEventTypes.CHANGE:
start_pos = body_stream.tell()
changes = {}
while body_stream.tell() < start_pos + so_body_size:
attrib_name = decoder.readString()
attrib_value = decoder.readElement()
assert attrib_name not in changes, (attrib_name,changes.keys())
changes[attrib_name] = attrib_value
assert body_stream.tell() == start_pos + so_body_size,\
(body_stream.tell(),start_pos,so_body_size)
event['data'] = changes
elif event['type'] == SOEventTypes.MESSAGE:
start_pos = body_stream.tell()
msg_params = []
while body_stream.tell() < start_pos + so_body_size:
msg_params.append(decoder.readElement())
assert body_stream.tell() == start_pos + so_body_size,\
(body_stream.tell(),start_pos,so_body_size)
event['data'] = msg_params
elif event['type'] == SOEventTypes.CLEAR:
assert so_body_size == 0, so_body_size
event['data'] = ''
elif event['type'] == SOEventTypes.DELETE:
event['data'] = decoder.readString()
elif event['type'] == SOEventTypes.USE_SUCCESS:
assert so_body_size == 0, so_body_size
event['data'] = ''
else:
assert False, event['type']
return event
class RtmpWriter:
""" This class writes RTMP messages into a stream. """
chunk_size = 128
def __init__(self, stream):
"""
Initialize the RTMP writer and set it to write into the specified
stream.
"""
self.stream = stream
def flush(self):
""" Flush the underlying stream. """
self.stream.flush()
def write(self, message):
logging.debug('send %r', message)
""" Encode and write the specified message into the stream. """
datatype = message['msg']
body_stream = pyamf.util.BufferedByteStream()
encoder = pyamf.amf0.Encoder(body_stream)
if datatype == DataTypes.USER_CONTROL:
body_stream.write_ushort(message['event_type'])
body_stream.write(message['event_data'])
elif datatype == DataTypes.WINDOW_ACK_SIZE:
body_stream.write_ulong(message['window_ack_size'])
elif datatype == DataTypes.SET_PEER_BANDWIDTH:
body_stream.write_ulong(message['window_ack_size'])
body_stream.write_uchar(message['limit_type'])
elif datatype == DataTypes.COMMAND:
for command in message['command']:
encoder.writeElement(command)
elif datatype == DataTypes.SHARED_OBJECT:
encoder.serialiseString(message['obj_name'])
body_stream.write_ulong(message['curr_version'])
body_stream.write(message['flags'])
for event in message['events']:
self.write_shared_object_event(event, body_stream)
else:
assert False, message
self.send_msg(datatype, body_stream.getvalue())
def write_shared_object_event(self, event, body_stream):
"""
Helper method that writes one shared object inside a shared object RTMP
message.
"""
inner_stream = pyamf.util.BufferedByteStream()
encoder = pyamf.amf0.Encoder(inner_stream)
event_type = event['type']
if event_type == SOEventTypes.USE:
assert event['data'] == '', event['data']
elif event_type == SOEventTypes.CHANGE:
for attrib_name in event['data']:
attrib_value = event['data'][attrib_name]
encoder.serialiseString(attrib_name)
encoder.writeElement(attrib_value)
elif event['type'] == SOEventTypes.CLEAR:
assert event['data'] == '', event['data']
elif event['type'] == SOEventTypes.USE_SUCCESS:
assert event['data'] == '', event['data']
else:
assert False, event
body_stream.write_uchar(event_type)
body_stream.write_ulong(len(inner_stream))
body_stream.write(inner_stream.getvalue())
def send_msg(self, datatype, body):
"""
Helper method that send the specified message into the stream. Takes
care to prepend the necessary headers and split the message into
appropriately sized chunks.
"""
# Values that just work. :-)
if datatype >= 1 and datatype <= 7:
channel_id = 2
stream_id = 0
else:
channel_id = 3
stream_id = 0
timestamp = 0
header = rtmp_protocol_base.Header(
channelId=channel_id,
streamId=stream_id,
datatype=datatype,
bodyLength=len(body),
timestamp=timestamp)
rtmp_protocol_base.header_encode(self.stream, header)
for i in xrange(0,len(body),self.chunk_size):
chunk = body[i:i+self.chunk_size]
self.stream.write(chunk)
if i+self.chunk_size < len(body):
rtmp_protocol_base.header_encode(self.stream, header, header)
class FlashSharedObject:
"""
This class represents a Flash Remote Shared Object. Its data are located
inside the self.data dictionary.
"""
def __init__(self, name):
"""
Initialize a new Flash Remote SO with a given name and empty data.
"""
self.name = name
self.data = {}
self.use_success = False
def use(self, reader, writer):
"""
Initialize usage of the SO by contacting the Flash Media Server. Any
remote changes to the SO should be now propagated to the client.
"""
self.use_success = False
msg = {
'msg': DataTypes.SHARED_OBJECT,
'curr_version': 0,
'flags': '\x00\x00\x00\x00\x00\x00\x00\x00',
'events': [
{
'data': '',
'type': SOEventTypes.USE
}
],
'obj_name': self.name
}
writer.write(msg)
writer.flush()
def handle_message(self, message):
"""
Handle an incoming RTMP message. Check if it is of any relevance for the
specific SO and process it, otherwise ignore it.
"""
if message['msg'] == DataTypes.SHARED_OBJECT and \
message['obj_name'] == self.name:
events = message['events']
if not self.use_success:
assert events[0]['type'] == SOEventTypes.USE_SUCCESS, events[0]
assert events[1]['type'] == SOEventTypes.CLEAR, events[1]
events = events[2:]
self.use_success = True
self.handle_events(events)
return True
else:
return False
def handle_events(self, events):
""" Handle SO events that target the specific SO. """
for event in events:
event_type = event['type']
if event_type == SOEventTypes.CHANGE:
for key in event['data']:
self.data[key] = event['data'][key]
self.on_change(key)
elif event_type == SOEventTypes.DELETE:
key = event['data']
assert key in self.data, (key,self.data.keys())
del self.data[key]
self.on_delete(key)
elif event_type == SOEventTypes.MESSAGE:
self.on_message(event['data'])
else:
assert False, event
def on_change(self, key):
pass
def on_delete(self, key):
pass
def on_message(self, data):
pass
class RtmpClient:
""" Represents an RTMP client. """
def __init__(self, ip, port, tc_url, page_url, swf_url, app):
""" Initialize a new RTMP client. """
self.ip = ip
self.port = port
self.tc_url = tc_url
self.page_url = page_url
self.swf_url = swf_url
self.app = app
self.shared_objects = []
def handshake(self):
""" Perform the handshake sequence with the server. """
self.stream.write_uchar(3)
c1 = rtmp_protocol_base.Packet()
c1.first = 0
c1.second = 0
c1.payload = 'x'*1528
c1.encode(self.stream)
self.stream.flush()
self.stream.read_uchar()
s1 = rtmp_protocol_base.Packet()
s1.decode(self.stream)
c2 = rtmp_protocol_base.Packet()
c2.first = s1.first
c2.second = s1.second
c2.payload = s1.payload
c2.encode(self.stream)
self.stream.flush()
s2 = rtmp_protocol_base.Packet()
s2.decode(self.stream)
def connect_rtmp(self, connect_params):
""" Initiate a NetConnection with a Flash Media Server. """
msg = {
'msg': DataTypes.COMMAND,
'command':
[
u'connect',
1,
{
'videoCodecs': 252,
'audioCodecs': 3191,
'flashVer': u'WIN 10,1,85,3',
'app': self.app,
'tcUrl': self.tc_url,
'videoFunction': 1,
'capabilities': 239,
'pageUrl': self.page_url,
'fpad': False,
'swfUrl': self.swf_url,
'objectEncoding': 0
}
]
}
msg['command'].extend(connect_params)
self.writer.write(msg)
self.writer.flush()
while True:
msg = self.reader.next()
if self.handle_message_pre_connect(msg):
break
def call(self, proc_name, parameters = {}, trans_id = 0):
""" Runs remote procedure calls (RPC) at the receiving end. """
msg = {
'msg': DataTypes.COMMAND,
'command':
[
proc_name,
trans_id,
parameters
]
}
self.writer.write(msg)
self.writer.flush()
def handle_message_pre_connect(self, msg):
""" Handle messages arriving before the connection is established. """
if msg['msg'] == DataTypes.COMMAND:
assert msg['command'][0] == '_result', msg
assert msg['command'][1] == 1, msg
assert msg['command'][3]['code'] == \
'NetConnection.Connect.Success', msg
return True
elif msg['msg'] == DataTypes.WINDOW_ACK_SIZE:
assert msg['window_ack_size'] == 2500000, msg
elif msg['msg'] == DataTypes.SET_PEER_BANDWIDTH:
assert msg['window_ack_size'] == 2500000, msg
assert msg['limit_type'] == 2, msg
elif msg['msg'] == DataTypes.USER_CONTROL:
assert msg['event_type'] == UserControlTypes.STREAM_BEGIN, msg
assert msg['event_data'] == '\x00\x00\x00\x00', msg
elif msg['msg'] == DataTypes.SET_CHUNK_SIZE:
assert msg['chunk_size'] > 0 and msg['chunk_size'] <= 65536, msg
self.reader.chunk_size = msg['chunk_size']
else:
assert False, msg
return False
def connect(self, connect_params):
""" Connect to the server with the given connect parameters. """
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.socket.connect((self.ip, self.port))
self.file = self.socket.makefile()
self.stream = FileDataTypeMixIn(self.file)
self.handshake()
self.reader = RtmpReader(self.stream)
self.writer = RtmpWriter(self.stream)
self.connect_rtmp(connect_params)
def shared_object_use(self, so):
""" Use a shared object and add it to the managed list of SOs. """
if so in self.shared_objects:
return
so.use(self.reader, self.writer)
self.shared_objects.append(so)
def handle_messages(self):
""" Start the message handling loop. """
while True:
msg = self.reader.next()
handled = self.handle_simple_message(msg)
if handled:
continue
for so in self.shared_objects:
if so.handle_message(msg):
handled = True
break
if not handled:
assert False, msg
def handle_simple_message(self, msg):
""" Handle simple messages, e.g. ping requests. """
if msg['msg'] == DataTypes.USER_CONTROL and msg['event_type'] == \
UserControlTypes.PING_REQUEST:
resp = {
'msg':DataTypes.USER_CONTROL,
'event_type':UserControlTypes.PING_RESPONSE,
'event_data':msg['event_data'],
}
self.writer.write(resp)
self.writer.flush()
return True
return False