From d630d58096b550092f222d0489f4dd8aa34ffe10 Mon Sep 17 00:00:00 2001
From: samwiseg0 <samwiseg00@gmail.com>
Date: Mon, 6 Aug 2018 04:11:56 -0400
Subject: [PATCH] Add options and pull counts

---
 ombi.py | 100 +++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 74 insertions(+), 26 deletions(-)

diff --git a/ombi.py b/ombi.py
index 60c88fa..5aa1812 100644
--- a/ombi.py
+++ b/ombi.py
@@ -1,39 +1,87 @@
 # Do not edit this script. Edit configuration.py
+import sys
 import requests
 from datetime import datetime, timezone
 from influxdb import InfluxDBClient
-
+import argparse
+from argparse import RawTextHelpFormatter
 import configuration
 
-current_time = datetime.now(timezone.utc).astimezone().isoformat()
-
 headers = {'Apikey': configuration.ombi_api_key}
-get_tv_requests = requests.get('{}/api/v1/Request/tv'.format(configuration.ombi_url), headers=headers).json()
-get_movie_requests = requests.get('{}/api/v1/Request/movie'.format(configuration.ombi_url), headers=headers).json()
 
-count_movie_requests = 0
-count_tv_requests = 0
+def now_iso():
+    now_iso = datetime.now(timezone.utc).astimezone().isoformat()
+    return now_iso
 
-for show in get_tv_requests:
-    count_tv_requests += 1
-    
-for movie in get_movie_requests:
-    count_movie_requests += 1
+def influx_sender(influx_payload):
+    influx = InfluxDBClient(configuration.influxdb_url, configuration.influxdb_port, configuration.influxdb_username,
+                            configuration.influxdb_password, configuration.ombi_influxdb_db_name)
+    influx.write_points(influx_payload)
 
-influx_payload = [
-    {
-        "measurement": "Ombi",
-        "tags": {
-            "type": "Requests"
-        },
-        "time": current_time,
-        "fields": {
-            "total": count_movie_requests + count_tv_requests
+def get_total_requests():
+    get_tv_requests = requests.get('{}/api/v1/Request/tv'.format(configuration.ombi_url), headers=headers).json()
+    get_movie_requests = requests.get('{}/api/v1/Request/movie'.format(configuration.ombi_url), headers=headers).json()
+
+    count_movie_requests = 0
+    count_tv_requests = 0
+
+    for show in get_tv_requests:
+        count_tv_requests += 1
+
+    for movie in get_movie_requests:
+        count_movie_requests += 1
+
+    influx_payload = [
+        {
+            "measurement": "Ombi",
+            "tags": {
+                "type": "Request_Total"
+            },
+            "time": now_iso(),
+            "fields": {
+                "total": count_movie_requests + count_tv_requests
+            }
         }
-    }
-]
+    ]
+    return influx_payload
 
-influx = InfluxDBClient(configuration.grafana_url, configuration.grafana_port, configuration.grafana_username,
-                        configuration.grafana_password, configuration.ombi_grafana_db_name)
-influx.write_points(influx_payload)
+def get_request_counts():
+    get_request_counts = requests.get('{}/api/v1/Request/count'.format(configuration.ombi_url), headers=headers).json()
 
+    influx_payload = [
+        {
+        "measurement": "Ombi",
+            "tags": {
+                "type": "Request_Counts"
+            },
+            "time": now_iso(),
+            "fields": {
+                "pending": int(get_request_counts['pending']),
+                "approved": int(get_request_counts['approved']),
+                "available": int(get_request_counts['available'])
+            }
+        }
+    ]
+    return influx_payload
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(prog='Ombi stats operations',
+        description='Script to aid in data gathering from Ombi', formatter_class=RawTextHelpFormatter)
+
+    parser.add_argument("--total",  action='store_true',
+        help='Get the total count of all requests')
+
+    parser.add_argument("--counts",  action='store_true',
+        help='Get the count of pending, approved, and available requests')
+
+    opts = parser.parse_args()
+
+    if opts.total:
+        influx_sender(get_total_requests())
+
+    elif opts.counts:
+        influx_sender(get_request_counts())
+
+    elif len(sys.argv) == 1:
+        parser.print_help(sys.stderr)
+        sys.exit(1)