bhdupload/bhdupload.py

#!/usr/bin/env python

import argparse
import atexit
import guessit
import json
import math
import requests
import os
import subprocess
import sys
import tempfile
from dotenv import load_dotenv


def absoluteFilePaths(directory):
    """ Get full file paths within a specific directory

    Keyword arguments:
    directory -- directory path containing files
    """
    files = []
    for dirpath, _, filenames in os.walk(directory):
        for f in sorted(filenames):
            files.append(os.path.abspath(os.path.join(dirpath, f)))
    return files


def readFiles(files, key='image[]'):
    """ Open files for Requests payload

    Keyword arguments:
    files -- list of files (full paths)
    key -- Requests payload form field name
    """
    result = []
    for f in files:
        result.append((key, (os.path.basename(f), open(f, 'rb'))))
    return result


def getArgs():
    """ Get command-line arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('input', action='store', type=str, help='Input file')
    parser.add_argument('-m', dest='screenshot', action='store', type=int, default=1200, help='Time to screenshot')
    parser.add_argument('-c', dest='count', action='store', type=int, default=4, help='Amount of screenshots to take')
    parser.add_argument('-u', dest='images', action='store_false', default=True,
                        help='Do not Create & Upload Screenshots')
    parser.add_argument('-t', dest='torrent', action='store_false', default=True,
                        help='Do not Create Torrent File')
    parser.add_argument('-l', dest='live', action='store_true', help='Upload instead of going to drafts')
    parser.add_argument('-a', dest='anon', action='store_true', help='Set an Anonymous upload')
    parser.add_argument('-v', dest='debug', action='store_true', help='Enable extra debug mode')
    return parser.parse_args()


def getDuration(file):
    """ Get duration of video file (in seconds)

    Keyword arguments:
    file -- input file (full path)
    """
    p = subprocess.Popen([
        'ffprobe',
        '-v', 'error',
        '-loglevel', '0',
        '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1',
        file
    ], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    duration = p.stdout.read().decode('utf-8')
    retval = p.wait()
    if retval == 127:
        raise ValueError('ffprobe is not installed.')
    return int(math.floor(float(duration)))


def getPieceSize(file):
    """ Calculates a respectible piece size
    """
    size = os.path.getsize(file)
    mb = size/1024/1024
    if mb > 72000:
        piece = 23
    if mb < 72000:
        piece = 22
    if mb < 16000:
        piece = 21
    if mb < 8000:
        piece = 20
    return piece


def mkScreenshot(file, outputDir, offset):
    """ Make a screenshot from file in specified output dir, with the specified offset

    Keyword arguments:
    file -- input file (full path)
    outputDir -- directory to write screenshot to
    offset -- how many seconds into the file to pull the screenshot from
    """
    p = subprocess.Popen([
        'ffmpeg',
        '-ss', str(offset),
        '-r', '1',
        '-i', file,
        outputDir + '/ss' + str(offset).zfill(5) + '.png'
    ], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    retval = p.wait()
    if retval == 127:
        raise ValueError('ffmpeg is not installed.')
    return True


def mkScreenshots(file, count, tempdir, interval):
    """ Make screenshots from file, with the specified interval between screenshots

    Keyword arguments:
    file -- input file (full path)
    count -- amount of screenshots to take
    tempdir -- our temporary directory
    interval -- interval between screenshots (seconds)
    """

    for i in range(1, count + 1):
        sys.stdout.write(str(i) + '.. ')
        sys.stdout.flush()
        mkScreenshot(file, tempdir, i * interval)
    print('Done!')
    return absoluteFilePaths(tempdir)


def uploadScreenshots(files):
    """ Upload screenshots to Imgbox

    Keyword arguments:
    files - list of images (full paths)
    """
    newargs = ['imgbox']
    for img in files:
        newargs.append(img)
    newargs.append("-w")
    newargs.append("350")
    newargs.append("-j")

    # print(f"Running: {newargs} to create screenshots")
    p = subprocess.run(newargs, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if p.check_returncode():
        print(f"Possible issue uploading img: {p.stderr}")
    return p.stdout


def removeDir(directory):
    """ Removes a directory including all files in that directory

    Keyword arguments:
    directory -- directory path containing files
    """
    files = absoluteFilePaths(directory)
    for f in files:
        os.remove(f)
    os.rmdir(directory)


def mktorrent(passkey, directory, tempdir):
    piecelength = getPieceSize(directory)
    base = f"mktorrent -p -l {piecelength} -a https://beyond-hd.me/announce/{passkey}"
    output = f"{tempdir}/torrent.torrent"

    args = base.split()
    args.append("-o")
    args.append(output)
    args.append(directory)

    p = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if p.check_returncode():
        print(f"{directory} failed to create torrent file: {output}")
        return "failure"
    return f"{output}"


def mediainfo(directory, tempdir):
    args = ["mediainfo", f"{directory}"]
    mediafile = open(f"{tempdir}/mediainfo.txt", 'w')

    p = subprocess.run(args, stdout=mediafile)
    if p.check_returncode():
        print(f"{directory} failed to create mediainfo file")
        mediafile.flush()
        return "failure"
    mediafile.flush()
    return f"{tempdir}/mediainfo.txt"


def find_info(name, page=1):
    torrent_name = str(name)
    print(f"Searching for {torrent_name}...")
    # first search on name
    if page == 1:
        r = SESS.post(BHD_SEARCH, data={'action': 'search', 'search': torrent_name})
    else:
        r = SESS.post(BHD_SEARCH, data={'action': 'search', 'search': torrent_name, 'page': page})
    try:
        if r.status_code == 200:
            json = r.json()
            if json["total_results"] >= 1:
                if args.debug:
                    print(f"Total Results: {json['total_results']}")
                print(f"Current Page: {page} of {json['total_pages']}")
                for item in json["results"]:
                    guess = dict(guessit.guessit(item['name']))
                    title = (guess['title'])
                    if args.debug:
                        print(f"Current Item: Title: {title} - {torrent_name}")
                    if title == torrent_name:
                        # We found an exact match
                        imdb = item["imdb_id"]
                        tmdb = item["tmdb_id"]
                        category = item["category"]
                        pack = item["tv_pack"]
                        print(f"Found: {torrent_name} IMDB: {imdb}, TMDB: {tmdb}, Category: {category}, TV?: {pack} ")
                        return imdb, tmdb, category, pack
                # If we get to the end of a page we need to check there are not more pages
                while page != json["total_pages"]:
                    next_page = page + 1
                    imdb, tmdb, category, pack = find_info(name, next_page)
                    return imdb, tmdb, category, pack
                # If we get to here we've never found {torrent_name}
                print(f"Failed to find {torrent_name}")
                return False, False, False, False
        return imdb, tmdb, category, pack

    except Exception as e:
        print(f"Error finding torrent information: {e}")
        print(f"{json}")
    return "", "", "", ""


def find_gaps(name, _source, _reso=""):
    '''
    payload = {'name': name,
               'description': BBCODE,
               'category_id': category,
               'type': reso,
               'source': source,
               'imdb_id': imdbid,
               'tmdb_id': tmdbid,
               'pack': pack,
               'live': live,
               'anon': anon}

    '''
    torrent_name = str(name)
    source = str(_source)
    reso = str(_reso)
    action = ""
    print(f"Searching for {torrent_name} - {source} - {reso}...")
    # first search on name
    r = SESS.post(BHD_SEARCH, data={
        'action': 'search',
        'search': torrent_name,
        'source': source,
        'types': reso})
    try:
        if r.status_code == 200:
            _json = r.json()
            if _json['total_results'] >= 1:
                action = False
                print(f"Found {_json['total_results']} matching our criteria...")
                for result in _json['results']:
                    if result['seeders'] <= 1:
                        print(f"{result['name']} - Could be dead")
                        action = True
                    elif result['seeders'] >= 1 and result['seeders'] <= 3:
                        print(f"{result['name']} - Needs saving, not uploading")
                        action = False
            else:
                # No torrents match our Source & Resolution, We can upload safely!
                print(f"Couldn't find torrents matching: {torrent_name}.{source}.{reso}")
                action = True
            return action
    except Exception as e:
        print(f"Error finding torrent information: {e}")
        print(f"{_json}")
    return False


# MAIN
# initialize variables
load_dotenv()

BHD_PASSKEY = os.getenv("BHD_PASSKEY")
BHD_API = os.getenv("BHD_API")
BHD_SEARCH = f'https://beyond-hd.me/api/torrents/{BHD_API}'
BHD_UPLOAD = f'https://beyond-hd.me/api/upload/{BHD_API}'
SESS = requests.session()
BBCODE = ""
tempdir = tempfile.mkdtemp()
atexit.register(removeDir, tempdir)

args = getArgs()

directory = args.input

if args.images:
    # Get duration of input file
    duration = getDuration(args.input)
    # count number of screenshots to generate
    count = int(args.count)
    interval = args.screenshot/count

    print(f'We are going to generate {count} @ {interval}s')

    # Generate screenshots
    print('Generating screenshots...')
    screenshots = mkScreenshots(args.input, count, tempdir, interval)
    if len(screenshots) == 0:
        raise ValueError('No screenshots generated.  Quitting!')
    print(str(len(screenshots)) + ' screenshots generated.')

    # Upload screenshots
    print('Uploading screenshots...')
    uploads = uploadScreenshots(screenshots)

# Check result for successful uploads
    BBCODE = str()
    uploaded = json.loads(uploads)
    for up in uploaded:
        if up["success"] is not True:
            continue
        url = up["web_url"]
        img = up["thumbnail_url"]

        BBCODE += f"[URL={url}][IMG]{img}[/IMG][/URL]"
    BBCODE += "\n"

# Output results
    print("BBCode:\n")
    print(BBCODE)

if args.torrent:
    file_name = os.path.basename(args.input)

    print("Guessing information...")
    # USE GUESSIT TO DEFINE MORE VARIABLES
    guess = dict(guessit.guessit(args.input))
    searchterm = (guess['title'])
    reso = guess.get('screen_size', None)
    source = (guess['source'])
    other = guess.get('other', None)
    lastchar = file_name[len(file_name)-4:]
    if lastchar == ".mkv" or lastchar == ".avi" or lastchar == ".mp4":
        name = str(file_name.replace(".", " ")[:-4])
    else:
        name = str(file_name.replace(".", " "))

    # JANKY IF STATEMENTS TO MAKE GUESSIT OUTPUT API FRIENDLY
    if "Web" in source:
        source = "WEB"
    elif ((source == "Blu-ray") and (other == "Remux") and (reso == "1080p")):
        reso = "BD Remux"
    elif "DVD" in name:
        reso = "DVD Remux"
    elif ((source == "Ultra HD Blu-ray") and (other == "Remux") and (reso == "2160p")):
        reso = "UHD Remux"
        source = "Blu-ray"

    # Find the required info before we upload (hopefully BHD already has it uploaded)
    imdbid, tmdbid, category, pack = find_info(searchterm)

    live = 0
    anon = 0
    if args.live:
        live = 1
    if args.anon:
        anon = 1
    if category == "TV":
        category = 2
    if category == "Movies":
        category = 1

    payload = {'name': name,
               'description': BBCODE,
               'category_id': category,
               'type': reso,
               'source': source,
               'imdb_id': imdbid,
               'tmdb_id': tmdbid,
               'pack': pack,
               'live': live,
               'anon': anon}

    # Notify if it failed to be found
    if imdbid is False or tmdbid is False or category is False:
        print(f"Failed to find {searchterm} already on BHD - You might be the first!",
              "You will need to edit it in drafts before uploading")
        if live == 1:
            print(f"We are sending {searchterm} to drafts instead of live. Fix it up!")
            live = 0

    # Generate MediaInfo
    print("Creating mediainfo...")
    mediainfo_file = mediainfo(args.input, tempdir)
    # Generate TorrentFile
    print("Creating torrent...")
    torrent_file = mktorrent(BHD_PASSKEY, args.input, tempdir)

    if args.debug:
        print("")
        print(f"Debug: We are about to submit the following: {payload}")
        print("")
    torrentfile = f"{torrent_file}"
    mediainfofile = f"{mediainfo_file}"
    files = {'file': open(torrentfile, 'rb'), 'mediainfo': open(mediainfofile, 'rb')}

    response = requests.request("POST", BHD_UPLOAD, data=payload, files=files)
    print("")
    print(response.text)
# Cleanup
print('Done!')