plugins.ustreamtv: Refactor to support more streams.

1 files changed, 62 insertions, 24 deletions

diff --git a/src/livestreamer/plugins/ustreamtv.py b/src/livestreamer/plugins/ustreamtv.py
index 35e17f4..76fda5c 100644
--- a/src/livestreamer/plugins/ustreamtv.py
+++ b/src/livestreamer/plugins/ustreamtv.py
@@ -1,13 +1,16 @@
 from livestreamer.compat import str, bytes
+from livestreamer.packages.flashmedia import AMF0Packet, AMFError
 from livestreamer.plugins import Plugin, PluginError, NoStreamsError
-from livestreamer.stream import RTMPStream
+from livestreamer.stream import HLSStream, RTMPStream
 from livestreamer.utils import urlget
 
-import xml.dom.minidom, re
+from io import BytesIO
+
+import re
 
 class UStreamTV(Plugin):
     AMFURL = "http://cgw.ustream.tv/Viewer/getStream/1/{0}.amf"
-    SWFURL = "http://cdn1.ustream.tv/swf/4/viewer.rsl.210.swf"
+    SWFURL = "http://static-cdn1.ustream.tv/swf/live/viewer3:50.swf"
 
     @classmethod
     def can_handle_url(self, url):
@@ -20,35 +23,70 @@ class UStreamTV(Plugin):
         if match:
             return int(match.group(1))
 
-    def _get_streams(self):
-        def get_amf_value(data, key):
-            pattern = ("{0}\x02..(.*?)\x00").format(key)
-            match = re.search(bytes(pattern, "ascii"), data)
-            if match:
-                return str(match.group(1), "ascii")
+    def _create_stream(self, cdn, streamname):
+        url = "{0}/{1}".format(cdn, streamname)
+        options = dict(rtmp=url, pageUrl=self.url,
+                       swfUrl=self.SWFURL, live=True)
+        return RTMPStream(self.session, options)
 
-        streams = {}
+    def _get_streams(self):
         channelid = self._get_channel_id(self.url)
 
         if not channelid:
             raise NoStreamsError(self.url)
 
+
         self.logger.debug("Fetching stream info")
         res = urlget(self.AMFURL.format(channelid))
-        data = res.content
-
-        playpath = get_amf_value(data, "streamName")
-        cdnurl = get_amf_value(data, "cdnUrl")
-        fmsurl = get_amf_value(data, "fmsUrl")
-
-        if playpath:
-            stream = RTMPStream(self.session, {
-                "rtmp": ("{0}/{1}").format(cdnurl or fmsurl, playpath),
-                "pageUrl": self.url,
-                "swfUrl": self.SWFURL,
-                "live": True
-            })
-            streams["live"] = stream
+
+        try:
+            packet = AMF0Packet.deserialize(BytesIO(res.content))
+        except (IOError, AMFError) as err:
+            raise PluginError(("Failed to parse AMF packet: {0}").format(str(err)))
+
+        result = None
+        for message in packet.messages:
+            if message.target_uri == "/1/onResult":
+                result = message.value
+                break
+
+        if not result:
+            raise PluginError("No result found in AMF packet")
+
+        streams = {}
+
+        if "liveHttpUrl" in result:
+            try:
+                hlsstreams = HLSStream.parse_variant_playlist(self.session,
+                                                              result["liveHttpUrl"])
+                streams.update(hlsstreams)
+            except IOError as err:
+                self.logger.warning("Failed to get variant playlist: {0}", err)
+
+        if "streamName" in result:
+            if "cdnUrl" in result:
+                cdn = result["cdnUrl"]
+            elif "fmsUrl" in result:
+                cdn = result["fmsUrl"]
+            else:
+                self.logger.warning("Missing cdnUrl and fmsUrl from result")
+                return streams
+
+            if "videoCodec" in result and result["videoCodec"]["height"] > 0:
+                streamname = "{0}p".format(int(result["videoCodec"]["height"]))
+            else:
+                streamname = "live"
+
+            streams[streamname] = self._create_stream(cdn, result["streamName"])
+
+        if "streamVersions" in result:
+            for version, info in result["streamVersions"].items():
+                if "streamVersionCdn" in info:
+                    for name, cdn in info["streamVersionCdn"].items():
+                        if "cdnStreamUrl" in cdn and "cdnStreamName" in cdn:
+                            streams["cdn_" + name] = self._create_stream(cdn["cdnStreamUrl"],
+                                                                         cdn["cdnStreamName"])
+
 
         return streams