make_nfo.py

import os
import sys
import json
from datetime import datetime
import log

from stash_interface import StashInterface
import config

def basename(f):
    f = os.path.normpath(f)
    return os.path.basename(f)

def xmlSafe(text):
    if text:
        return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    return text

def getOutputSTRMFile(scene):
    if config.filename == "filename":
        return os.path.join(config.save_path, "{}.".format(os.path.splitext(basename(scene["path"]))[0]) + config.playlist_ext)
    elif config.filename == "stashid":
        sceneID = scene["id"]
        return os.path.join(config.save_path, "{}.".format(sceneID) + config.playlist_ext)
    else:
        log.error("Check your config file.")

def getOutputNFOFile(scene):
    if config.filename == "filename":
        return os.path.join(config.save_path, "{}.nfo".format(os.path.splitext(basename(scene["path"]))[0]))
    elif config.save_path == "with files":
        return os.path.join(os.path.dirname(scene["path"]), "{}.nfo".format(os.path.splitext(basename(scene["path"]))[0]))
    elif config.filename == "stashid":
        sceneID = scene["id"]
        return os.path.join(config.save_path, "{}.nfo".format(sceneID))
    else:
        log.error("Check your config file.")

def getSceneTitle(scene):
    if scene["title"] != None and scene["title"] != "":
        return scene["title"]
    return basename(scene["path"])

def getGenreTags():
    query = """
    {
      findTags(
        tag_filter: {name: {value: """ + '"' + config.genre_parentname + '"' + """ , modifier: EQUALS}}
        filter: {per_page: -1}
      ) {
        count
        tags {
          id
          name
          children {
            name
            id
          }
        }
      }
    }
    """
    results = stash.graphql_query(query)
    if not results["findTags"]["tags"]:
        log.warning("'genre_parentname' tag not found in stash. Check README & config.")
        return []
    else:
        resultschildren = results["findTags"]["tags"][0]["children"]
        return get_ids(resultschildren)

def get_ids(obj):
    ids = []
    for item in obj:
        ids.append(item['id'])
    return ids

def URLrewrite(url):
    if config.replacement_server_url:
        url = url.replace(config.server_url, config.replacement_server_url)
    if config.api_key:
        return url + "&apikey=" + config.api_key
    return url

def generateSTRM(scene):
    if config.m3u:
        if config.playlist_ext == "m3u8":
            return "#EXTM3U\n#EXTENC: UTF-8\n#EXTINF:" + str(int(scene["file"]["duration"])) + "," + getSceneTitle(scene) + "\n" + URLrewrite(scene["paths"]["stream"])
        else:
            return "#EXTM3U\n#EXTINF:" + str(int(scene["file"]["duration"])) + "," + getSceneTitle(scene) + "\n" + URLrewrite(scene["paths"]["stream"])
    else:
        return URLrewrite(scene["paths"]["stream"])

def generateNFO(scene):
    ret = """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<!-- Generated by StashNfoExporterKodi plugin on {now} -->
<movie>
    <title>{title}</title>
    <uniqueid type="stash" default="true">{id}</uniqueid>
{stash_ids}
    <mpaa>XXX</mpaa>
    <playcount>{o_counter}</playcount>
    <dateadded>{created_at}</dateadded>
    <premiered>{date}</premiered>
    <userrating>{rating}</userrating>
    <plot>{details}</plot>
    <studio>{studio}</studio>
    <director>{director}</director>
{performers}
{set}
{tags}
{thumbs}
{fanart}
{genres}
</movie>"""

    allgenres = getGenreTags()
    genres = []

    stash_ids = ""
    for id in scene["stash_ids"]:
        stash_ids = stash_ids + """    <uniqueid type="stashdb" endpoint="{}">{}</uniqueid>\n""".format(id["endpoint"], id["stash_id"])

    tags = ""
    for t in scene["tags"]:
        tags = tags + """    <tag>{}</tag>\n""".format(xmlSafe(t["name"]))
        if t["id"] in allgenres:
            genres.append("    <genre>{}</genre>\n".format(xmlSafe(t["name"])))

    rating = ""
    if scene["rating"] != None:
        rating = scene["rating"]

    date = ""
    if scene["date"] != None:
        date = scene["date"]

    studio = ""
    logo = ""
    if scene["studio"] != None:
        studio = scene["studio"]["name"]
        logo = scene["studio"]["image_path"]
        if not logo.endswith("?default=true"):
            logo = URLrewrite(logo)
        else:
            logo = ""

    performers = ""
    i = 0
    for p in scene["performers"]:
        thumb = URLrewrite(p["image_path"])
        performers = performers + """    <actor>
        <name>{}</name>
        <role></role>
        <order>{}</order>
        <thumb aspect="poster">{}</thumb>
    </actor>\n""".format(xmlSafe(p["name"]), i, xmlSafe(thumb))
        i += 1

    thumbs = ["""    <thumb aspect="landscape">{}</thumb>""".format(xmlSafe(URLrewrite(scene["paths"]["screenshot"])))]

    set = ""
    director = ""
    gotposter = False

    # this code is really only for one movie:
    # if you have multiple movies, the director and movie data will be only the last one found.
    # Patch accepted if you need more for some weird reason
    if "movies" in scene:
        log.debug(scene["movies"])
        for mlist in scene["movies"]:
            m = mlist["movie"]
            set = """    <set>
        <name>{}</name>
        <index>{}</index>
    </set>""".format(xmlSafe(m["name"]), mlist["scene_index"] or "")
            if "director" in m:
                director = m["director"]
            if "front_image_path" in m:
                poster = m["front_image_path"]
                if poster is not None:
                    thumbs.append("""    <thumb aspect="poster">{}</thumb>""".format(xmlSafe(URLrewrite(poster))))
                    gotposter = True
            if "back_image_path" in m:
                backcover = m["back_image_path"]
                if backcover:
                    thumbs.append("""    <thumb aspect="poster">{}</thumb>""".format(xmlSafe(URLrewrite(backcover))))

    # no movie poster, so let's (sadly) include the landscape screenshot as the poster, until we can do better.
    if gotposter == False:
        thumbs.append("""    <thumb aspect="poster">{}</thumb>""".format(xmlSafe(URLrewrite(scene["paths"]["screenshot"]))))

    fanart = ["""    <thumb aspect="fanart">{}</thumb>""".format(xmlSafe(URLrewrite(scene["paths"]["screenshot"])))]

    if logo != "":
        thumbs.append("""    <thumb aspect="clearlogo">{}</thumb>""".format(xmlSafe(URLrewrite(logo))))
        fanart.append("""    <thumb aspect="clearlogo">{}</thumb>""".format(xmlSafe(URLrewrite(logo))))

    fanart = """    <fanart>\n    {}\n    </fanart>""".format("\n".join(fanart))

    ret = ret.format(title = xmlSafe(getSceneTitle(scene)),
                     now = datetime.now(),
                     id = scene["id"],
                     endpoint = scene["id"],
                     stash_ids = stash_ids,
                     o_counter = scene["o_counter"] or "",
                     created_at = scene["created_at"][:10],
                     details = xmlSafe(scene["details"] or ""),
                     rating = rating,
                     date = date,
                     studio = xmlSafe(studio),
                     director = xmlSafe(director),
                     performers = performers,
                     set = set,
                     tags = tags,
                     thumbs = "\n".join(thumbs),
                     fanart = fanart,
                     genres = "\n".join(genres)
                     )
    ret = ret.replace("\n\n", "\n")
    return ret

def writeFile(fn, data, useUTF=False):
    encoding = None
    if useUTF:
        encoding = "utf-8-sig"
    os.makedirs(os.path.dirname(fn), exist_ok=True)
    f = open(fn, "w", encoding=encoding)
    f.write(data)
    f.close()

def main():
    global stash
    if len(sys.argv) > 1:
        # Loads from argv for testing...
        json_input = json.loads(sys.argv[1])
    else:
        json_input = json.loads(sys.stdin.read())
    stash = StashInterface(json_input["server_connection"])
    sceneID = json_input['args']['hookContext']['id']
    scene = stash.getScene(sceneID)

    # For existing nfo, config defines wether to proceed with generation or not...
    nfofilename = getOutputNFOFile(scene)
    if (not scene["organized"] and config.generate_when == "organized") or \
       (os.path.exists(nfofilename) and config.generate_when == "new"):
        # Skip generation...
        sys.exit()

    useUTF = True
    nfo = generateNFO(scene)
    writeFile(nfofilename, nfo, useUTF)

    if config.save_path != "with files":
        useUTF = False
        # don't write utf8 if it's strm extension, it breaks in Kodi, only use if m3u8
        if config.playlist_ext == "m3u8":
            useUTF = True
        strm = generateSTRM(scene)
        strmfilename = getOutputSTRMFile(scene)
        writeFile(strmfilename, strm, useUTF)

main()