diff --git a/killstream/kill_stream.py b/killstream/kill_stream.py index ebd40f1..69cc819 100644 --- a/killstream/kill_stream.py +++ b/killstream/kill_stream.py @@ -46,11 +46,13 @@ Tautulli > Settings > Notification Agents > New Script > Script Arguments: import os import sys -import argparse import json import time +import argparse from datetime import datetime -import requests +from requests import Session +from requests.adapters import HTTPAdapter +from requests.exceptions import RequestException TAUTULLI_URL = '' @@ -109,13 +111,15 @@ def debug_dump_vars(): .rjust(len(TAUTULLI_APIKEY), "x")) -def get_all_streams(user_id=None): +def get_all_streams(tautulli, user_id=None): """Get a list of all current streams. Parameters ---------- user_id : int The ID of the user to grab sessions for. + tautulli : obj + Tautulli object. Returns ------- objects @@ -131,17 +135,17 @@ def get_all_streams(user_id=None): return streams -def notify(opts, message, kill_type=None, stream=None): +def notify(all_opts, message, kill_type=None, stream=None, tautulli=None): """Decides which notifier type to use""" - if opts.notify and opts.richMessage: - rich_notify(opts.notify, opts.richMessage, opts.richColor, kill_type, - opts.serverName, opts.plexUrl, opts.posterUrl, message, stream) - elif opts.notify: - basic_notify(opts.notify, opts.sessionId, opts.username, message) + if all_opts.notify and all_opts.richMessage: + rich_notify(all_opts.notify, all_opts.richMessage, all_opts.richColor, kill_type, + all_opts.serverName, all_opts.plexUrl, all_opts.posterUrl, message, stream, tautulli) + elif all_opts.notify: + basic_notify(all_opts.notify, all_opts.sessionId, all_opts.username, message, stream, tautulli) def rich_notify(notifier_id, rich_type, color=None, kill_type=None, server_name=None, - plex_url=None, poster_url=None, message=None, stream=None): + plex_url=None, poster_url=None, message=None, stream=None, tautulli=None): """Decides which rich notifier type to use. Set default values for empty variables Parameters @@ -150,6 +154,8 @@ def rich_notify(notifier_id, rich_type, color=None, kill_type=None, server_name= The ID of the user to grab sessions for. rich_type : str Contains 'discord' or 'slack'. + color : Union[int, str] + Hex string or integer representation of color. kill_type : str The kill type used. server_name : str @@ -162,9 +168,13 @@ def rich_notify(notifier_id, rich_type, color=None, kill_type=None, server_name= Message sent to the client. stream : obj Stream object. + tautulli : obj + Tautulli object. """ notification = Notification(notifier_id, SUBJECT_TEXT, BODY_TEXT, tautulli, stream) - + # Initialize Variables + title = '' + footer = '' # Set a default server_name if none is provided if server_name is None: server_name = 'Plex Server' @@ -207,7 +217,7 @@ def rich_notify(notifier_id, rich_type, color=None, kill_type=None, server_name= notification.send_slack(title, color, poster_url, plex_url, message, footer) -def basic_notify(notifier_id, session_id, username=None, message=None): +def basic_notify(notifier_id, session_id, username=None, message=None, stream=None, tautulli=None): """Basic notifier""" notification = Notification(notifier_id, SUBJECT_TEXT, BODY_TEXT, tautulli, stream) @@ -225,11 +235,11 @@ class Tautulli: self.apikey = apikey self.debug = debug - self.session = requests.Session() - self.adapters = requests.adapters.HTTPAdapter(max_retries=3, - pool_connections=1, - pool_maxsize=1, - pool_block=True) + self.session = Session() + self.adapters = HTTPAdapter(max_retries=3, + pool_connections=1, + pool_maxsize=1, + pool_block=True) self.session.mount('http://', self.adapters) self.session.mount('https://', self.adapters) @@ -240,14 +250,14 @@ class Tautulli: import urllib3 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) - def _call_api(self, cmd, payload, method='GET', debug=None): + def _call_api(self, cmd, payload, method='GET'): payload['cmd'] = cmd payload['apikey'] = self.apikey try: response = self.session.request(method, self.url + '/api/v2', params=payload) - except: - print("Tautulli request failed for cmd '{}'. Invalid Tautulli URL?".format(cmd)) + except RequestException as e: + print("Tautulli request failed for cmd '{}'. Invalid Tautulli URL? Error: {}".format(cmd, e)) if self.debug: traceback.print_exc() return @@ -303,11 +313,11 @@ class Tautulli: class Stream: def __init__(self, session_id=None, user_id=None, username=None, tautulli=None, session=None): + self.state = None self.session_id = session_id self.user_id = user_id self.username = username self.session_exists = False - self.tautulli = tautulli if session is not None: @@ -585,32 +595,32 @@ if __name__ == "__main__": debug_dump_vars() # Create a Tautulli instance - tautulli = Tautulli(TAUTULLI_URL.rstrip('/'), TAUTULLI_APIKEY, VERIFY_SSL, opts.debug) + tautulli_server = Tautulli(TAUTULLI_URL.rstrip('/'), TAUTULLI_APIKEY, VERIFY_SSL, opts.debug) # Create initial Stream object with basic info - stream = Stream(opts.sessionId, opts.userId, opts.username, tautulli) + tautulli_stream = Stream(opts.sessionId, opts.userId, opts.username, tautulli_server) # Only pull all stream info if using richMessage if opts.notify and opts.richMessage: - stream.get_all_stream_info() + tautulli_stream.get_all_stream_info() # Set a default message if none is provided if opts.killMessage: - message = ' '.join(opts.killMessage) + kill_message = ' '.join(opts.killMessage) else: - message = 'The server owner has ended the stream.' + kill_message = 'The server owner has ended the stream.' if opts.jbop == 'stream': - stream.terminate(message) - notify(opts, message, 'Stream', stream) + tautulli_stream.terminate(kill_message) + notify(opts, kill_message, 'Stream', tautulli_stream, tautulli_server) elif opts.jbop == 'allStreams': - streams = get_all_streams(opts.userId) - for stream in streams: - tautulli.terminate_session(session_id=stream.session_id, message=message) - notify(opts, message, 'All Streams', stream) + all_streams = get_all_streams(tautulli_server, opts.userId) + for a_stream in all_streams: + tautulli_server.terminate_session(session_id=a_stream.session_id, message=kill_message) + notify(opts, kill_message, 'All Streams', a_stream, tautulli_server) elif opts.jbop == 'paused': - killed_stream = stream.terminate_long_pause(message, opts.limit, opts.interval) + killed_stream = tautulli_stream.terminate_long_pause(kill_message, opts.limit, opts.interval) if killed_stream: - notify(opts, message, 'Paused', stream) + notify(opts, kill_message, 'Paused', tautulli_stream, tautulli_server) diff --git a/requirements.txt b/requirements.txt index 1bd2be7..374efd7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ # pip install -r requirements.txt #--------------------------------------------------------- requests -plexapi \ No newline at end of file +plexapi +urllib3 \ No newline at end of file