Overhaul + Sonarr support!

6 files changed, 735 insertions, 0 deletions

diff --git a/code/utils/config.py b/code/utils/config.py
new file mode 100644
index 0000000..491677d
--- /dev/null
+++ b/code/utils/config.py
@@ -0,0 +1,298 @@
+import jsonschema
+import validators
+import yaml
+import sys
+import logging
+import requests
+import sqlite3
+from colorlog import ColoredFormatter
+
+
+log_level = logging.DEBUG
+log_format = (
+    "  %(log_color)s%(levelname)-8s%(reset)s |"
+    " %(log_color)s%(message)s%(reset)s"
+)
+
+logging.root.setLevel(log_level)
+formatter = ColoredFormatter(log_format)
+
+stream = logging.StreamHandler()
+stream.setLevel(log_level)
+stream.setFormatter(formatter)
+
+LOG = logging.getLogger("pythonConfig")
+LOG.setLevel(log_level)
+LOG.addHandler(stream)
+
+BOT_TOKEN = None
+
+RADARR_ENABLED = False
+RADARR_HOST_URL = None
+RADARR_HEADERS = None
+RADARR_ROOT_FOLDER_PATH = None
+RADARR_QUALITY_PROFILE_ID = None
+
+SONARR_ENABLED = False
+SONARR_HOST_URL = None
+SONARR_HEADERS = None
+SONARR_ROOT_FOLDER_PATH = None
+SONARR_QUALITY_PROFILE_ID = None
+
+JELLYFIN_ENABLED = False
+JELLYFIN_URL = None
+JELLYFIN_HEADERS = None
+ACCOUNT_TIME = None
+SIMPLE_PASSWORDS = False
+
+schema = {
+    "type": "object",
+    "properties": {
+        "bot_info": {
+            "type": "object",
+            "properties": {
+                "bot_token": {"type": "string"},
+            },
+            "required": ["bot_token"],
+        },
+        "radarr": {
+            "type": "object",
+            "properties": {
+                "host_url": {"type": "string"},
+                "api_key": {"type": "string"},
+                "root_folder_path": {"type": "string"},
+                "quality_profile_id": {"type": "integer"},
+            },
+            "required": [
+                "host_url",
+                "api_key",
+                "root_folder_path",
+            ],
+        },
+        "sonarr": {
+            "type": "object",
+            "properties": {
+                "host_url": {"type": "string"},
+                "api_key": {"type": "string"},
+                "root_folder_path": {"type": "string"},
+                "quality_profile_id": {"type": "integer"},
+            },
+            "required": [
+                "host_url",
+                "api_key",
+                "root_folder_path",
+            ],
+        },
+        "jellyfin": {
+            "type": "object",
+            "properties": {
+                "url": {"type": "string"},
+                "api_key": {"type": "string"},
+                "account_time": {"type": "integer"},
+            },
+            "required": ["url", "api_key", "account_time"],
+        },
+    },
+    "required": ["bot_info", "radarr", "sonarr"],
+}
+
+
+def load_config() -> None:
+    """
+    Load DB, then load and validate the config file
+    If the file does not exist, generate it
+    """
+    database_setup()
+    try:
+        with open("config.yaml", "r") as f:
+            contents = f.read()
+            validate_config(contents)
+
+    except FileNotFoundError:
+        with open("config.yaml", "w") as f:
+            f.write(
+                """
+bot_info:
+    bot_token: YOUR_BOT_TOKEN
+
+radarr:
+    host_url: RADARR_URL
+    api_key: RADARR_API_KEY
+    root_folder_path: RADARR_ROOT_FOLDER_PATH
+    quality_profile_id: RADARR_QUALITY_PROFILE_ID
+
+sonarr:
+    host_url: SONARR_URL
+    api_key: SONARR_API_KEY
+    root_folder_path: SONARR_ROOT_FOLDER_PATH
+    quality_profile_id: SONARR_QUALITY_PROFILE_ID
+
+jellyfin:
+    url: JELLYFIN_URL
+    api_key: JELLYFIN_API_KEY
+    account_time: ACCOUNT_ACTIVE_TIME
+    simple_passwords: SIMPLE_OR_COMPLEX_PASSWORDS
+                """
+            )
+
+        sys.exit(
+            LOG.critical(
+                "Config file `config.yaml` has been generated. Input necessary"
+                " fields and restart. Refer to README for help!"
+            )
+        )
+
+
+def database_setup() -> None:
+    """
+    Create the database if it does not exist
+    """
+    db = sqlite3.connect("cordarr.db")
+    cursor = db.cursor()
+    cursor.execute(
+        "CREATE TABLE IF NOT EXISTS requests (title TEXT, release_year TEXT,"
+        " local_id INTEGER, tmdbid INTEGER, tvdbid INTEGER, user_id INTEGER)"
+    )
+    cursor.execute(
+        "CREATE TABLE IF NOT EXISTS jellyfin_accounts (user_id INTEGER,"
+        " jellyfin_user_id INTEGER, deletion_time DATETIME)"
+    )
+    db.commit()
+    db.close()
+
+
+def validate_config(contents) -> None:
+    """
+    Validate the contents of the config file and assign variables
+
+    Args:
+        contents (str): The contents of the config file
+    """
+    global BOT_TOKEN, RADARR_HOST_URL, RADARR_ENABLED, RADARR_HEADERS, RADARR_ROOT_FOLDER_PATH, RADARR_QUALITY_PROFILE_ID, SONARR_ENABLED, SONARR_HOST_URL, SONARR_HEADERS, SONARR_ROOT_FOLDER_PATH, SONARR_QUALITY_PROFILE_ID, JELLYFIN_ENABLED, JELLYFIN_URL, JELLYFIN_HEADERS, ACCOUNT_TIME, SIMPLE_PASSWORDS
+
+    config = yaml.safe_load(contents)
+
+    try:
+        jsonschema.validate(config, schema)
+    except jsonschema.ValidationError as e:
+        sys.exit(LOG.critical(f"Error in config.yaml: {e.message}"))
+
+    #
+    # Begin validating values and assigning variables
+    #
+
+    BOT_TOKEN = config["bot_info"]["bot_token"]
+
+    if "radarr" in config:
+        if not validators.url(config["radarr"]["host_url"]):
+            sys.exit(
+                LOG.critical(
+                    "Error in config.yaml: Invalid URL for Radarr host"
+                )
+            )
+        else:
+            RADARR_HOST_URL = config["radarr"]["host_url"]
+
+        RADARR_HEADERS = {
+            "Content-Type": "application/json",
+            "X-Api-Key": config["radarr"]["api_key"],
+        }
+        RADARR_ROOT_FOLDER_PATH = config["radarr"]["root_folder_path"]
+        # set radarr quality profile id
+        RADARR_QUALITY_PROFILE_ID = validate_profile(
+            "radarr", RADARR_HOST_URL, RADARR_HEADERS, config
+        )
+        RADARR_ENABLED = True
+
+    if "sonarr" in config:
+        if not validators.url(config["sonarr"]["host_url"]):
+            sys.exit(
+                LOG.critical(
+                    "Error in config.yaml: Invalid URL for Sonarr host"
+                )
+            )
+        else:
+            SONARR_HOST_URL = config["sonarr"]["host_url"]
+
+        SONARR_HEADERS = {
+            "Content-Type": "application/json",
+            "X-Api-Key": config["sonarr"]["api_key"],
+        }
+        SONARR_ROOT_FOLDER_PATH = config["sonarr"]["root_folder_path"]
+        # set sonarr quality profile id
+        SONARR_QUALITY_PROFILE_ID = validate_profile(
+            "sonarr", SONARR_HOST_URL, SONARR_HEADERS, config
+        )
+        SONARR_ENABLED = True
+
+    if "jellyfin" in config:
+        if not validators.url(config["jellyfin"]["url"]):
+            LOG.critical(
+                "Error in config.yaml: Invalid URL for Jellyfin - account"
+                " creation disabled"
+            )
+        else:
+            JELLYFIN_URL = config["jellyfin"]["url"]
+
+        JELLYFIN_HEADERS = {
+            "Content-Type": "application/json",
+            "X-Emby-Token": config["jellyfin"]["api_key"],
+        }
+        ACCOUNT_TIME = config["jellyfin"]["account_time"]
+        SIMPLE_PASSWORDS = config["jellyfin"]
+        JELLYFIN_ENABLED = True
+
+
+def validate_profile(
+    service: str, url: str, headers: dict, config: dict
+) -> int:
+    """
+    Validate the quality profile ID for the given service
+
+    Args:
+        service (str): The service to validate the profile for
+        url (str): The URL of the service
+        headers (dict): The headers for the request
+        config (dict): The config file
+
+    Returns:
+        int: The quality profile ID
+    """
+    profiles = requests.get(f"{url}/api/v3/qualityProfile", headers=headers)
+
+    if profiles.status_code != 200:
+        LOG.critical(
+            f"Error in config.yaml: Unable to get {service} quality profiles."
+            f" API Key invalid or incorrect {service} URL"
+        )
+
+    else:
+        profiles = profiles.json()
+        # If ID is not given, list options
+        if "quality_profile_id" not in config[f"{service}"]:
+            LOG.critical(
+                "Error in config.yaml: No quality profile ID provided for"
+                f" {service}. Look below for a list of your available"
+                " profiles:"
+            )
+
+            for profile in profiles:
+                LOG.info(f"ID: {profile['id']} | Name: {profile['name']}")
+
+        # ID is given, validate
+        else:
+            quality_profile_id = config[f"{service}"]["quality_profile_id"]
+            if quality_profile_id not in [
+                profile["id"] for profile in profiles
+            ]:
+                LOG.critical(
+                    f"Error in config.yaml: Invalid {service} quality profile"
+                    " ID. Look below for your available profiles:"
+                )
+
+                for profile in profiles:
+                    LOG.info(f"ID: {profile['id']} | Name: {profile['name']}")
+                sys.exit()
+            # Everything valid, assign
+            else:
+                return quality_profile_id
diff --git a/code/utils/content_add.py b/code/utils/content_add.py
new file mode 100644
index 0000000..fda43af
--- /dev/null
+++ b/code/utils/content_add.py
@@ -0,0 +1,57 @@
+import requests
+
+
+def add_content(
+    content_info: dict,
+    service: str,
+    host: str,
+    headers: str,
+    folder_path: str,
+    profile_id: str,
+):
+    """
+    Add content to Sonarr or Radarr
+
+    Args:
+        content_info (dict): The content information
+        service (str): The service to add the content to
+        host (str): The host URL
+        headers (str): The headers for the request
+        folder_path (str): The folder path to download the content to
+        profile_id (str): The profile ID to download the content in
+
+    Returns:
+        str: The ID of the content or False
+    """
+    # Get the content data based on ID
+    data = requests.get(
+        url=(
+            f"{host}/api/v3/movie/lookup/tmdb?tmdbId={content_info['contentId']}"
+            if service == "radarr"
+            else f"{host}/api/v3/series/lookup?term=tvdb:{content_info['contentId']}"
+        ),
+        headers=headers,
+    ).json()[0]
+
+    data["monitored"] = True
+    data["qualityProfileId"] = profile_id
+    data["rootFolderPath"] = folder_path
+    # Search for the content on add
+    data["addOptions"] = {
+        (
+            "searchForMovie"
+            if service == "radarr"
+            else "searchForMissingEpisodes"
+        ): True
+    }
+    # Send the request to add the content
+    response = requests.post(
+        f"{host}/api/v3/{'movie' if service == 'radarr' else 'series'}",
+        headers=headers,
+        json=data,
+    )
+
+    if response.status_code == 201:
+        return response.json()["id"]
+    else:
+        return False
diff --git a/code/utils/content_get.py b/code/utils/content_get.py
new file mode 100644
index 0000000..1b724e5
--- /dev/null
+++ b/code/utils/content_get.py
@@ -0,0 +1,59 @@
+import requests
+
+
+def get_content(
+    query: str,
+    service: str,
+    host: str,
+    headers: str,
+):
+    """
+    Fetch the top 5 results from the service given a query
+
+    Args:
+        query (str): The query to search for
+        service (str): The service to search in
+        host (str): The host URL
+        headers (str): The headers for the request
+
+    Returns:
+        list: A list containing content_info dict
+        str: NO RESULTS
+        str: ALREADY ADDED
+    """
+    query = query.strip().replace(" ", "%20")
+    # Search for matching content
+    results = requests.get(
+        f"{host}/api/v3/{'movie' if service == 'radarr' else 'series'}/lookup?term={query}",
+        headers=headers,
+    ).json()
+
+    if len(results) == 0:
+        return "NO RESULTS"
+    # If already added to library
+    if results[0]["added"] != "0001-01-01T05:51:00Z":
+        return "ALREADY ADDED"
+
+    # Add info for top results
+    content_info = []
+    for i in range(min(5, len(results))):
+        content_info.append(
+            {
+                "title": results[i]["title"],
+                "year": results[i]["year"],
+                "contentId": results[i][
+                    f"{'tmdbId' if service == 'radarr' else 'tvdbId'}"
+                ],
+                "description": results[i]["overview"],
+            }
+        )
+
+        # Add remotePoster field, set None if not available
+        try:
+            content_info[i]["remotePoster"] = results[i]["images"][0][
+                "remoteUrl"
+            ]
+        except IndexError:
+            content_info[i]["remotePoster"] = None
+
+    return content_info
diff --git a/code/utils/content_view.py b/code/utils/content_view.py
new file mode 100644
index 0000000..e8df29c
--- /dev/null
+++ b/code/utils/content_view.py
@@ -0,0 +1,203 @@
+import discord
+import sqlite3
+
+from utils.content_add import add_content
+
+"""
+View to add the Dropdown menu
+"""
+
+
+class AddContentView(discord.ui.View):
+    def __init__(
+        self,
+        content_data: list,
+        service: str,
+        host: str,
+        header: str,
+        path: str,
+        profile: int,
+        *,
+        timeout=180.0,
+    ):
+        super().__init__(timeout=timeout)
+        # Add the dropdown
+        self.add_item(
+            AddContentDropdown(
+                content_data, service, host, header, path, profile
+            )
+        )
+
+
+"""
+Dropdown containing the top 5 content results
+"""
+
+
+class AddContentDropdown(discord.ui.Select):
+    def __init__(
+        self,
+        content_data: list,
+        service: str,
+        host: str,
+        header: str,
+        path: str,
+        profile: str,
+        *,
+        timeout=180.0,
+    ):
+        self.content_data = content_data
+        self.service = service
+        self.host = host
+        self.header = header
+        self.path = path
+        self.profile = profile
+        options = []
+        for i in range(len(content_data)):
+            options.append(
+                discord.SelectOption(
+                    label=(
+                        f"{content_data[i]['title']} ({content_data[i]['year']})"
+                    ),
+                    description=f"Relevant ID: {content_data[i]['contentId']}",
+                    value=str(i),
+                )
+            )
+
+        super().__init__(
+            placeholder="Select from the dropdown",
+            options=options,
+        )
+
+    # Once an option has been selected
+    async def callback(self, interaction: discord.Interaction):
+        # Index of selected option
+        index = int(self.values[0])
+
+        # Add selected contents info to an embed
+        embed = discord.Embed(
+            title="Is the the content you want to add?",
+            description=(
+                f"**Title**: {self.content_data[index]['title']} | "
+                f"**Year**: {self.content_data[index]['year']}\n\n"
+                f"**Description**: {self.content_data[index]['description']}"
+            ),
+            color=0xD01B86,
+        )
+        embed.set_image(url=self.content_data[index]["remotePoster"])
+        # Change the view to the Request/Don't Request buttons
+        view = RequestButtonView(
+            self.content_data[index],
+            self.service,
+            self.host,
+            self.header,
+            self.path,
+            self.profile,
+        )
+        await interaction.response.edit_message(embed=embed, view=view)
+
+
+"""
+View containing the "Request" and "Don't Request" buttons
+"""
+
+
+class RequestButtonView(discord.ui.View):
+    def __init__(
+        self,
+        content_info: dict,
+        service: str,
+        host: str,
+        headers: str,
+        path: str,
+        profile: int,
+        *,
+        timeout=180.0,
+    ):
+        super().__init__(timeout=timeout)
+        self.content_info = content_info
+        self.service = service
+        self.host = host
+        self.headers = headers
+        self.path = path
+        self.profile = profile
+
+    @discord.ui.button(label="Request", style=discord.ButtonStyle.success)
+    async def request_button(
+        self, interaction: discord.Interaction, button: discord.ui.Button
+    ):
+        # Add the content to the relevant library
+        local_id = add_content(
+            self.content_info,
+            self.service,
+            self.host,
+            self.headers,
+            self.path,
+            self.profile,
+        )
+
+        # Alert the user that the content has been added
+        if local_id:
+            embed = discord.Embed(
+                title="Content Added",
+                description=(
+                    f"**{self.content_info['title']}** has been added to the"
+                    f" {self.service} library. Check the status of your"
+                    " requested content with `/status`."
+                ),
+                color=0xD01B86,
+            )
+            await interaction.response.send_message(embed=embed)
+        # Alert the user that the content failed to be added
+        else:
+            embed = discord.Embed(
+                title="Failed to Add Content",
+                description=(
+                    "An error occured when attempting to add"
+                    f" **{self.content_info['title']}** to the"
+                    f" {self.service} library."
+                ),
+            )
+            return await interaction.response.send_message(embed=embed)
+
+        # Keep track of the requests for the `/status` command
+        db = sqlite3.connect("cordarr.db")
+        cursor = db.cursor()
+        cursor.execute(
+            "INSERT INTO requests (title, release_year, local_id, tmdbid,"
+            " tvdbid, user_id) VALUES (?, ?, ?, ?, ?, ?)",
+            (
+                self.content_info["title"],
+                self.content_info["year"],
+                local_id,
+                (
+                    self.content_info["contentId"]
+                    if self.service == "radarr"
+                    else None
+                ),
+                (
+                    None
+                    if self.service == "radarr"
+                    else self.content_info["contentId"]
+                ),
+                interaction.user.id,
+            ),
+        )
+        db.commit()
+        db.close()
+
+    @discord.ui.button(label="Don't Request", style=discord.ButtonStyle.danger)
+    async def dont_request_button(
+        self, interaction: discord.Interaction, button: discord.ui.Button
+    ):
+        embed = discord.Embed(
+            title="Request Cancelled",
+            description=(
+                "Request has been cancelled. If you would like to request a"
+                " different"
+                f" {'movie' if self.service == 'radarr' else 'show'}, run the"
+                " `/request` command again."
+            ),
+            color=0xD01B86,
+        )
+        await interaction.response.send_message(embed=embed)
diff --git a/code/utils/jellyfin_create.py b/code/utils/jellyfin_create.py
new file mode 100644
index 0000000..08c5230
--- /dev/null
+++ b/code/utils/jellyfin_create.py
@@ -0,0 +1,80 @@
+import datetime
+import requests
+import random
+import sqlite3
+from wonderwords import RandomWord
+from string import ascii_lowercase, digits
+
+from utils.config import (
+    JELLYFIN_URL,
+    JELLYFIN_HEADERS,
+    ACCOUNT_TIME,
+    SIMPLE_PASSWORDS,
+)
+
+
+def create_jellyfin_account(user_id):
+    """
+    Create a new Jellyfin account for the user and return the username and password
+
+    Args:
+        user_id (int): Discord user ID to create the account for
+
+    Returns:
+        tuple: The username and password of the new Jellyfin account
+    """
+    # Create username/password
+    username = RandomWord().word(word_min_length=5, word_max_length=5)
+    if SIMPLE_PASSWORDS:
+        password = RandomWord().word(word_min_length=5, word_max_length=10)
+    else:
+        password = "".join(random.choices(ascii_lowercase + digits, k=15))
+
+    deletion_time = datetime.datetime.now() + datetime.timedelta(
+        minutes=ACCOUNT_TIME * 60
+    )
+    # Create the new Jellyfin account
+    request_1 = requests.post(
+        f"{JELLYFIN_URL}/Users/New",
+        headers=JELLYFIN_HEADERS,
+        json={"Name": username, "Password": password},
+    )
+
+    if request_1.status_code != 200:
+        return False
+
+    # Get the user ID of the new account
+    jellyfin_user_id = request_1.json()["Id"]
+    # Get the account policy and make edits
+    request_2 = requests.get(
+        f"{JELLYFIN_URL}/Users/{jellyfin_user_id}", headers=JELLYFIN_HEADERS
+    )
+    if request_2.status_code != 200:
+        return False
+
+    account_policy = request_2.json()
+    account_policy["Policy"]["SyncPlayAccess"] = "JoinGroups"
+    account_policy["Policy"]["EnableContentDownloading"] = False
+    account_policy["Policy"]["InvalidLoginAttemptCount"] = 3
+    account_policy["Policy"]["MaxActiveSessions"] = 1
+    # Update the user with the newly edited policy
+    request_3 = requests.post(
+        f"{JELLYFIN_URL}/Users?userId={jellyfin_user_id}",
+        headers=JELLYFIN_HEADERS,
+        json=account_policy,
+    )
+    if request_3.status_code != 204:
+        return False
+
+    # Add the information to the database
+    db = sqlite3.connect("cordarr.db")
+    cursor = db.cursor()
+    cursor.execute(
+        "INSERT INTO jellyfin_accounts (user_id, jellyfin_user_id,"
+        " deletion_time) VALUES (?, ?, ?)",
+        (user_id, jellyfin_user_id, deletion_time),
+    )
+    db.commit()
+    db.close()
+
+    return username, password
diff --git a/code/utils/jellyfin_delete.py b/code/utils/jellyfin_delete.py
new file mode 100644
index 0000000..6164e40
--- /dev/null
+++ b/code/utils/jellyfin_delete.py
@@ -0,0 +1,38 @@
+import datetime
+import sqlite3
+import requests
+
+from utils.config import JELLYFIN_URL, JELLYFIN_HEADERS
+
+
+def delete_accounts():
+    """
+    Delete Jellyfin accounts that have passed their deletion time
+    """
+    # Get all expired Jellyfin accounts
+    db = sqlite3.connect("cordarr.db")
+    cursor = db.cursor()
+    cursor.execute(
+        "SELECT jellyfin_user_id FROM jellyfin_accounts WHERE"
+        " deletion_time < ?",
+        (datetime.datetime.now(),),
+    )
+    jellyfin_user_ids = cursor.fetchall()
+
+    # Delete the Jellyfin accounts
+    for jellyfin_user_id in jellyfin_user_ids:
+        request = requests.delete(
+            f"{JELLYFIN_URL}/Users/{jellyfin_user_id[0]}",
+            headers=JELLYFIN_HEADERS,
+        )
+        # If 204 - account deleted
+        # If 404 - account not found
+        # Either way, remove account from database
+        if request.status_code in (404, 204):
+            cursor.execute(
+                "DELETE FROM jellyfin_accounts WHERE jellyfin_user_id = ?",
+                (jellyfin_user_id,),
+            )
+
+    db.commit()
+    db.close()