import argparse import os import requests import math import re import json import config BATCH_SIZE = 100 def parseArgs(): parser = argparse.ArgumentParser(description="Generate nfo and strm files for Kodi integration.") parser.add_argument("mode", metavar="MODE", choices=["generate-nfo", "generate-strm"], help="generate-nfo or generate-strm") parser.add_argument("--inline", action="store_true", help="Generate nfo files along side video files") parser.add_argument("--outdir", metavar="", help="Generate files in ") parser.add_argument("--preserve-path", action="store_true", help="Include source file directory structure in output directory (with --outdir only)") parser.add_argument("--truncate-prefix", type=str, metavar="", help="Remove prefix from output directory (with --preserve-path only)") parser.add_argument("--use-source-filenames", action="store_true", help="Use source filenames for strm files instead of stash id") parser.add_argument("--overwrite", action="store_true", help="Overwrite nfo/strm files if already present") parser.add_argument("--filter", metavar="", help="JSON graphql string to filter scenes with") parser.add_argument("--genre", metavar="", help="Genre to assign. May be included multiple times", action="append") return parser.parse_args() # raw plugins may accept the plugin input from stdin, or they can elect # to ignore it entirely. In this case it optionally reads from the # command-line parameters. def main(): args = parseArgs() if args.mode == "generate-nfo": generateNFOFiles(args) elif args.mode == "generate-strm": generateSTRMFiles(args) def generateNFOFiles(args): if not args.inline and args.outdir == "": print("--outdir or --inline must be set\n") return filter = args.filter or "" if filter != "": filter = json.loads(filter) else: filter = {} total = getCount(filter) pages = math.ceil(total / BATCH_SIZE) i = 1 while i <= pages: print("Processing page {} of {}".format(i, pages)) scenes = getScenes(i, filter) for scene in scenes: # don't regenerate if file already exists and not overwriting output = getOutputNFOFile(scene["path"], args) if not args.overwrite and os.path.exists(output): continue nfo = generateNFO(scene, args) writeFile(output, nfo, True) i += 1 def generateSTRMFiles(args): if args.outdir == "": print("--outdir must be set\n") return filter = args.filter or "" if filter != "": filter = json.loads(filter) else: filter = {} total = getCount(filter) pages = math.ceil(total / BATCH_SIZE) i = 1 while i <= pages: print("Processing page {} of {}".format(i, pages)) scenes = getScenes(i, filter) for scene in scenes: name = "" outdir = getOutputDir(scene["path"], args) if args.use_source_filenames: name = basename(os.path.splitext(scene["path"])[0]) else: name = scene["id"] name = os.path.join(outdir, name) # don't regenerate if file already exists and not overwriting strmOut = name + ".strm" if args.overwrite or not os.path.exists(strmOut): data = generateSTRM(scene) writeFile(strmOut, data, False) output = name + ".nfo" if args.overwrite or not os.path.exists(output): nfo = generateNFO(scene, args) writeFile(output, nfo, True) i += 1 def basename(f): f = os.path.normpath(f) return os.path.basename(f) def getOutputSTRMFile(sceneID, args): return os.path.join(args.outdir, "{}.strm".format(sceneID)) def getOutputDir(sourceFile, args): ret = args.outdir if args.preserve_path: if args.truncate_prefix != None: toRemove = args.truncate_prefix if sourceFile.startswith(toRemove): sourceFile = sourceFile[len(toRemove):] sourceFile = os.path.normpath(sourceFile) ret = os.path.join(args.outdir, os.path.dirname(sourceFile)) return ret def getOutputNFOFile(sourceFile, args): if args.inline: # just replace the extension return os.path.splitext(sourceFile)[0]+".nfo" outdir = getOutputDir(sourceFile, args) ret = os.path.join(outdir, basename(sourceFile)) return os.path.splitext(ret)[0]+".nfo" def getCount(sceneFilter): query = """ query findScenes($filter: FindFilterType!, $scene_filter: SceneFilterType!) { findScenes(filter: $filter, scene_filter: $scene_filter) { count } } """ variables = { 'filter': { 'per_page': 0, }, 'scene_filter': sceneFilter } result = __callGraphQL(query, variables) return result["findScenes"]["count"] def getScenes(page, sceneFilter): query = """ query findScenes($filter: FindFilterType!, $scene_filter: SceneFilterType!) { findScenes(filter: $filter, scene_filter: $scene_filter) { scenes { id title path rating details date oshash paths { screenshot stream } studio { name image_path } performers { name image_path } tags { name } movies { movie { name } } } } } """ variables = { 'filter': { 'per_page': BATCH_SIZE, 'page': page, }, 'scene_filter': sceneFilter } result = __callGraphQL(query, variables) return result["findScenes"]["scenes"] def addAPIKey(url): if config.api_key: return url + "&apikey=" + config.api_key return url def getSceneTitle(scene): if scene["title"] != None and scene["title"] != "": return scene["title"] return basename(scene["path"]) def generateSTRM(scene): return scene["paths"]["stream"] def generateNFO(scene, args): ret = """ {title} {rating} {details} {id} {tags} {date} {studio} {performers} {thumbs} {fanart} {genres} """ tags = "" for t in scene["tags"]: tags = tags + """ {}""".format(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 = addAPIKey(logo) else: logo = "" performers = "" i = 0 for p in scene["performers"]: thumb = addAPIKey(p["image_path"]) performers = performers + """ {} {} {} """.format(p["name"], i, thumb) i += 1 thumbs = [ """{}""".format(addAPIKey(scene["paths"]["screenshot"])) ] fanart = [ """{}""".format(addAPIKey(scene["paths"]["screenshot"])) ] if logo != "": thumbs.append("""{}""".format(logo)) fanart.append("""{}""".format(logo)) fanart = """{}""".format("\n".join(fanart)) genres = [] if args.genre != None: for g in args.genre: genres.append("{}".format(g)) ret = ret.format(title = getSceneTitle(scene), rating = rating, id = scene["id"], tags = tags, date = date, studio = studio, performers = performers, details = scene["details"] or "", thumbs = "\n".join(thumbs), fanart = fanart, genres = "\n".join(genres)) return ret def writeFile(fn, data, useUTF): 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 __callGraphQL(query, variables = None): headers = { "Accept-Encoding": "gzip, deflate, br", "Content-Type": "application/json", "Accept": "application/json", "Connection": "keep-alive", "DNT": "1", "ApiKey": config.api_key } json = {} json['query'] = query if variables != None: json['variables'] = variables # handle cookies response = requests.post(config.server_url, json=json, headers=headers) if response.status_code == 200: result = response.json() if result.get("error", None): for error in result["error"]["errors"]: raise Exception("GraphQL error: {}".format(error)) if result.get("data", None): return result.get("data") else: raise Exception("GraphQL query failed:{} - {}. Query: {}. Variables: {}".format(response.status_code, response.content, query, variables)) main()