mirror of
https://github.com/stashapp/CommunityScripts.git
synced 2025-12-11 14:55:10 -06:00
335 lines
9.3 KiB
Python
335 lines
9.3 KiB
Python
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="<output directory>", help="Generate files in <outdir>")
|
|
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="<path prefix>", 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="<filter string>", help="JSON graphql string to filter scenes with")
|
|
parser.add_argument("--genre", metavar="<genre>", 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 = """
|
|
<?xml version="1.0" encoding="UTF-8" standalone="yes" ?>
|
|
<movie>
|
|
<title>{title}</title>
|
|
<userrating>{rating}</userrating>
|
|
<plot>{details}</plot>
|
|
<uniqueid type="stash">{id}</uniqueid>
|
|
{tags}
|
|
<premiered>{date}</premiered>
|
|
<studio>{studio}</studio>
|
|
{performers}
|
|
{thumbs}
|
|
{fanart}
|
|
{genres}
|
|
</movie>
|
|
"""
|
|
tags = ""
|
|
for t in scene["tags"]:
|
|
tags = tags + """
|
|
<tag>{}</tag>""".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 + """
|
|
<actor>
|
|
<name>{}</name>
|
|
<role></role>
|
|
<order>{}</order>
|
|
<thumb>{}</thumb>
|
|
</actor>""".format(p["name"], i, thumb)
|
|
i += 1
|
|
|
|
thumbs = [
|
|
"""<thumb aspect="poster">{}</thumb>""".format(addAPIKey(scene["paths"]["screenshot"]))
|
|
]
|
|
fanart = [
|
|
"""<thumb>{}</thumb>""".format(addAPIKey(scene["paths"]["screenshot"]))
|
|
]
|
|
if logo != "":
|
|
thumbs.append("""<thumb aspect="clearlogo">{}</thumb>""".format(logo))
|
|
fanart.append("""<thumb>{}</thumb>""".format(logo))
|
|
|
|
fanart = """<fanart>{}</fanart>""".format("\n".join(fanart))
|
|
|
|
genres = []
|
|
if args.genre != None:
|
|
for g in args.genre:
|
|
genres.append("<genre>{}</genre>".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() |