plugin.video.playstream/context_download.py

# -*- coding: utf-8 -*-
try:
    import wingdbstub
except:
    pass

import sys, os, urllib2, urllib, re, requests, datetime, time, json
#CLI_MODE = True
from kodiswift import xbmc, xbmcgui, CLI_MODE
from kodiswift import Plugin, storage
from resources.lib.content import util, ContentSources, file
from downloadqueue import DownloadQueue
import traceback


#from resources.lib.content import Downloader
#from twisted.web import client
#from twisted.internet import reactor, defer

plugin = Plugin(addon_id="plugin.video.playstream")
#plugin.load_addon_settings()
use_storage = plugin.get_setting("general_use_storage",bool) # TODO vajag nočekot vai nav labāk lietot pickle
overwrite = plugin.get_setting("general_download_overwrite",bool)
sleep_time = 5  # TODO jāliek iekš parametriem

cunicode = lambda s: s.decode("utf8") if isinstance(s, str) else s
cstr = lambda s: s.encode("utf8") if isinstance(s, unicode) else s

#print "argv=",sys.argv
cmd = sys.argv[1]
title = sys.argv[2]
data = sys.argv[3]
download_dir = sys.argv[4]
#overwrite = sys.argv[5] if len(sys.argv)>5 else False

queue_dir = os.path.join(xbmc.translatePath("special://temp"), "download_queue") if not CLI_MODE else "download_queue"
download_queue = DownloadQueue(queue_dir)

job_id = datetime.datetime.now().strftime("%Y%m%d%H%M")

if not os.path.exists(queue_dir):
    os.mkdir(queue_dir)


if cmd in ["download2", "download3"] and not CLI_MODE:
    download_dir2 = xbmcgui.Dialog().browse(0, "Select download folder", "files")
    if download_dir2:
        download_dir = download_dir2
    else:
        sys.exit()

#remote_dir = True if re.search("^\w+:", download_dir) else False
#print "encoding=",fs_encoding

cur_directory = os.path.dirname(__file__)
sources_directory = os.path.join(cur_directory,"resources","lib", "content", "sources")
sources = ContentSources.ContentSources(sources_directory)


def main():
    if not sources.is_video(data):  # Folderis
        n = 0
        dname = file.make_fname(title)
        download_dir2 = file.join(download_dir, dname)
        for current2 in sources.get_content(data):
            if sources.is_video(current2[1]):
                n += 1
                download_video(current2[0], current2[1], download_dir2, overwrite, num=n)
        #if n>0:
        #    notify("%s videos download started to %s"%(n,download_dir2))
        #else:
        #    notify("No videos to download")
    else:
        if cmd == "download3":
            dname = file.make_fname(title)
            download_dir2 = file.join(download_dir, dname)
        else:
            download_dir2 = download_dir
        download_video(title, data,download_dir2, overwrite)
        #self.msg("Video download started to %s"%download_dir)

def download_video(title, data, download_dir, overwrite, cb_notify=None, num=None):

    streams = sources.get_streams(data)
    if not streams:
        notify("No streams to download")
        return

    if len(streams)>1 and not num:
        slist = []
        for s in streams:
            slist.append("%s [%s,%s]"%(s["name"],s["quality"],s["lang"]))
        res = xbmcgui.Dialog().select("Select stream",slist) if not CLI_MODE else 0
        #res = xbmcgui.Dialog().contextmenu(slist) if not CLI_MODE else 0
        stream = streams[res]
    else:
        stream = streams[0]

    download_stream(stream, download_dir, overwrite, notify, num=num)

    #d = Downloader.download_video(stream["url"], os.path.join(download_dir, output), stream["headers"])
    #reactor.run()


    #xbmcgui.Dialog().ok("Info","Start download")

    #mode = "a" if os.path.exists("context_menu.log") else "w"
    #with open("context_menu.log", mode) as f:
    #    f.write("%s %s %s %s", sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])

def notify(text, title="Info", time=10000):
    if isinstance(text, unicode):
        text = text.encode("utf8")
    #xbmc.executebuiltin('Notification(Hello World,This is a simple example of notifications,5000,/script.hellow.world.png)')
    if CLI_MODE:
        print "Info: ", text
    else:
        xbmc.executebuiltin('Notification(%s, %s, %d, %s)'%("Info", text, time, xbmcgui.NOTIFICATION_INFO))

####################################################################################################
# Platform independent part, should be put in separate module
####################################################################################################

def download_stream(stream, download_dir, overwrite=True, cb_notify=None, num=False):
    #output =  stream["name"].replace("\\"," ").replace(":"," ").replace("|"," ")
    url = stream["url"]
    title = stream["name"].strip()
    output = file.make_fname(title)

    headers = stream["headers"] if "headers" in stream and stream["headers"] else {"user-agent":"Enigma2"}
    try:
        h = get_header(url,headers=headers)
        mtype = h.get("content-type")
        ext,stream_type = get_ext(mtype)
    except Exception as e:
        ext,stream_type = (".ts","hls")
    #stream_type = ContentSources.stream_type(url) #self.sources.stream_type(stream["url"])
    if not stream_type: #
        print "Not supported stream type found to download - %s"%(url)
        raise Exception("Not supported stream type found to download - %s"%(url))

    fname = output+ext
    outputfile = file.join(download_dir, fname)
    exists = file.exists(outputfile)
    if exists and num and num > 1:
        output = output + "_%02i" % num
        fname = output+ext
        outputfile = file.join(download_dir, fname)
        exists = file.exists(outputfile)

    if exists and not overwrite:
        if cb_notify: cb_notify("Download skipped - %s" % output)
        print "Download skipped - %s" % output
        return

    print "download_dir=", download_dir
    if not file.isdir(download_dir):
        try:
            file.mkdir(file.encode(download_dir))
        except Exception as e:
            traceback.print_exc()
            print 'Error creating download directory "%s"!\nPlease specify in the settings existing directory\n%s'%(download_dir,str(e))
            raise Exception('Error creating download directory "%s"!\nPlease specify in the settings existing directory\n%s'%(download_dir,str(e)))

    output_path = file.join(download_dir, output)

    if "nfo" in stream and stream["nfo"]:
        nfofile = file.join(download_dir, output+".nfo")
        #print "nfofile=", nfofile.encode("utf8")
        nfo_txt = util.nfo2xml(stream["nfo"])
        f = file.open(file.encode(nfofile),"w")
        f.write(nfo_txt)
        f.close()

    subfiles = []
    if "subs" in stream and stream["subs"]:
        for sub in stream["subs"]:
            suburl = sub["url"]
            slang = "." + sub["lang"] if sub["lang"] else ""
            download_sub(suburl, output+slang, download_dir)

    if "img" in stream and stream["img"]:
        download_image(stream["img"], output, download_dir)

    ### Start video file download ####
    job = {
        "job_id": job_id,
        "status": "downloading",
        "url":url,
        "output": output,
        "file": outputfile,
        "download_dir": download_dir,
        "headers": headers,
        "totalbytes": -1,
        "currentbytes": 0,
        "overwrite": overwrite,
    }
    download_queue.job_put(job_id, job)

    if cb_notify: cb_notify("%s put in download queue" % output)
    print "%s put in download queue" % output


    if stream_type == "hls":
        download_hls(url, outputfile, headers=headers, cb_notify=cb_notify)
    else:
        download_file(url, outputfile, headers=headers, cb_notify=cb_notify)


def download_hls(url, outputfile, headers=None, overwrite=True, limit=None, cb_notify=None):
    UA = "Mozilla/5.0 (Windows NT 10.0; WOW64; rv:51.0) Gecko/20100101 Firefox/51.0"
    if not headers:
        headers = {"User-Agent" : UA}
    key = headers["key"] if "key" in headers else ""
#    if not "User-Agent" in headers:
#        headers["User-Agent"] = UA
    #outputfile = file.join(download_dir,output)

    if cb_notify: cb_notify("Download started - %s" % outputfile)
    print "Download started - %s" % outputfile, url
    #print url


    try:
        r = requests.get(url,headers=headers)
    except Exception as e:
        raise Exception("Cannot open manifsest file - %s"%url)
    if not r.content.startswith("#EXTM3U"):
        raise Exception("Not valid manifest file - %s" % url)
    streams = re.findall(r"#EXT-X-STREAM-INF:.*?BANDWIDTH=(\d+).*?\n(.+?)$", r.content, re.IGNORECASE | re.MULTILINE)
    i = 0
    while streams:
        if i > 4: break
        sorted(streams, key=lambda item: int(item[0]), reverse=True)
        base_url = "/".join(url.split("?")[0].split("/")[:-1])+"/"
        url = streams[0][1]
        if not url.startswith("http"):
            url = base_url + url
        #print url
        try:
            r = requests.get(url, headers=headers)
        except Exception as e:
            raise Exception("Cannot open manifsest file - %s"%url)
        i += 1
        streams = re.findall(r"#EXT-X-STREAM-INF:.*?BANDWIDTH=(\d+).*?\n(.+?)$", r.content, re.IGNORECASE | re.MULTILINE)
    m = re.search('#EXT-X-KEY:METHOD=(.+?),URI="([^"]+)"', r.content, re.IGNORECASE)

    if m:
        url_key = m.group(2)
        r2 = requests.get(url_key, headers=headers)
        key = r2.content
    else:
        key = None

    ts_list = re.findall(r"#EXTINF:([\d\.]+),.*?\n(.+?)$", r.content, re.IGNORECASE | re.MULTILINE)
    base_url = "/".join(url.split("/")[:-1])+"/"

    if not len(ts_list):
        raise Exception("Cannot read fragment list in  manifsest file - %s"%url)

    ts_num = 0
    type = "vod" if "#EXT-X-ENDLIST" in r.content else "live"
    currentbytes = 0.0
    totalbytes = -1
    currenttime = 0.0
    totaltime = sum(map(float,zip(*ts_list)[0]))
    #ts_file = open(outputfile, "wb")

    tsfile = file.open(file.encode(outputfile),"wb")
    for ts in ts_list:
        url2 = ts[1]
        #print "Downloading ", url2
        #fname = os.path.join(download_dir,url2.split("/")[-1])
        if not url2.startswith("http"):
            url2 = base_url + url2

        r = requests.get(url2, headers=headers, verify=False)
        content = r.content
        if key:
            #from Crypto.Cipher import AES
            import pyaes
            iv = content[:16]
            #key2 = binascii.a2b_hex(key)
            #d = AES.new(key, AES.MODE_CBC, iv)
            d = pyaes.AESModeOfOperationCBC(key, iv = iv)
            content2 = ""
            for i in range(16, len(content), 16):
                content2 += d.decrypt(content[i:i+16])
            content = content2
            #content = d.decrypt(content[16:])
            #d = pyaes.Decrypter(pyaes.AESModeOfOperationCBC(key, iv))
            #content = d.feed(content[16:])
        tsfile.write(content)

        content_length = len(content)
        currentbytes += content_length
        currenttime += float(ts_list[ts_num][0])
        totalbytes = currentbytes * totaltime / currenttime

        # Checking queue
        while True:
            job = download_queue.job_get(job_id)
            if not job:
                print "Download canceled"
                notify("Download canceled - %s" % output)
                tsfile.close()
                os.remove(outputfile)
                return
            if job["status"] in ("waiting", "paused"):
                time.sleep(sleep_time)
                continue
            break
        job["currentbytes"] = currentbytes
        job["totalbytes"] = totalbytes
        download_queue.job_put(job_id, job)

        ts_num += 1
        #print "Fragment %s downloaded (%s)"%(self.ts_num,len(content))
        progress = float(currentbytes)/float(totalbytes)*100
        msg = "%.1f%% (%iMB/%iMB)"%(progress,currentbytes / 1024 / 1024,totalbytes / 1024 / 1024)
        #print msg
        #notify(msg)

        if type == "vod":
            if ts_num >= len(ts_list) or (limit and currenttime>limit):
                break
        else:  # live stream # TODO
            if limit and currenttime>limit: # TODO
                break

    download_queue.job_remove(job_id)
    print "Finished"
    notify("Download finished - %s" % outputfile)
    tsfile.close()


def download_file(url, outputfile, headers=None, overwrite=True, limit=None, cb_notify=None):
    global job
    UA = "Mozilla/5.0 (Windows NT 10.0; WOW64; rv:51.0) Gecko/20100101 Firefox/51.0"
    if not headers:
        headers = {"User-Agent" : UA}
    key = headers["key"] if "key" in headers else ""
#    if not "User-Agent" in headers:
#        headers["User-Agent"] = UA
    #fname = file.join(download_dir,output)

    if cb_notify: cb_notify("Download started - %s" % outputfile)
    print "Download started - %s" % outputfile
    #print url
    try:
        r = requests.get(url,headers=headers, stream=True)
    except Exception as e:
        raise Exception("Cannot open url - %s"%url)
    currentbytes = 0.0
    totalbytes = int(r.headers["content-length"])


    fd = file.open(file.encode(outputfile), 'wb')
    for chunk in r.iter_content(chunk_size=1024*1024):
        fd.write(chunk)
        currentbytes += len(chunk)
        progress = float(currentbytes)/float(totalbytes)*100
        msg = "%.1f%% (%iMB/%iMB)"%(progress,currentbytes / 1024 / 1024,totalbytes / 1024 / 1024)
        #print msg

        # Checking queue
        while True:
            job = download_queue.job_get(job_id)
            if not job:
                print "Download canceled"
                notify("Download canceled - %s" % outputfile)
                fd.close()
                os.remove(fname)
                return
            if job["status"] in ("waiting", "paused"):
                time.sleep(sleep_time)
                continue
            break
        job["currentbytes"] = currentbytes
        job["totalbytes"] = totalbytes
        download_queue.job_put(job_id, job)

        #notify(msg)

    fd.close()
    job.remove()
    print "Finished"
    notify("Download finished - %s" % outputfile)

def download_sub(suburl, output, download_dir):
    try:
        subs = urllib2.urlopen(suburl).read()
    except:
        subs = None
    if subs:
        if ".xml" in suburl:
            subs = util.ttaf2srt(subs)
            subext = ".srt"
        elif ".vtt" in suburl:
            subs = util.vtt2srt(subs)
            subext = ".srt"
        elif ".srt" in suburl:
            subext = ".srt"
        else:
            subext = ".srt"
        if subext:
            subfile = file.join(download_dir, output+subext)
            f = file.open(file.encode(subfile),"w")
            f.write(subs)
            f.close()
    else:
        print "\n Error downloading subtitle %s"%suburl
        raise Exception("Error downloading subtitle %s"%suburl)


def download_image(imgurl, output, download_dir):
    ext = imgurl.split("?")[0].split(".")[-1]
    ext = "." + ext
    ext = ".jpg"
    imgfile =  file.join(download_dir, output+ext)
    try:
        img = urllib2.urlopen(imgurl).read()
    except Exception(e):
        img = None
    if img:
        f = file.open(file.encode(imgfile),"wb")
        f.write(img)
        f.close()
    else:
        print "\n Error downloading image %s"%imgurl
        raise Exception("Error downloading image %s"%imgurl)



def get_header(url,headers=None):
    r = requests.head(url,headers=headers)
    return r.headers

def get_ext(mtype):
    stype = "http"
    if mtype in ("vnd.apple.mpegURL","application/x-mpegURL",'application/x-mpegurl',"application/vnd.apple.mpegurl"):
        return ".ts","hls"
    elif mtype in ("application/dash+xml"):
        return ".ts","dash" # TODO dash stream type  could be different !
    elif mtype in ("video/mp4"):
        return ".mp4","http"
    elif mtype in ("video/MP2T","video/mp2t"):
        return ".ts","http"
    elif mtype in ("video/x-flv"):
        return ".flv","http"
    elif mtype in ("video/quicktime"):
        return ".mov","http"
    elif mtype in ("video/x-msvideo"):
        return ".avi","http"
    elif mtype in ("video/x-ms-wmv"):
        return ".wmv","http"
    elif mtype in ("video/x-matroska"):
        return ".mkv","http"
    else:
        return ".mp4","http"


class Job(object):
    def __init__(self, job_id, queue_dir, job={}):
        self.job_id = job_id
        self.queue_dir = queue_dir
        self.job_file = os.path.join(queue_dir, job_id)
        self.job = job

    def read(self):
        try:
            with open(self.job_file, "r") as f:
                s = f.read()
            self.job = json.loads(s)
        except:
            self.job = {}
        return self.job

    def write(self):
        s = json.dumps(self.job)
        with open(self.job_file, "w") as f:
            f.write(s)
    def remove(self):
        if os.path.exists(self.job_file):
            os.remove(self.job_file)


if __name__ == '__main__':
    main()