add network_whitelist and rename whitelist to user_whitelist (#10)

AndrewPaglusch · Mar 4, 2023 · 8722a4f · 8722a4f
1 parent df0f581
commit 8722a4f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/app/config.ini.TEMPLATE b/app/config.ini.TEMPLATE
@@ -4,7 +4,8 @@ plex_url: ${PLEX_URL}
 plex_token: ${PLEX_TOKEN}
 ban_length_hrs: ${BAN_LENGTH_HRS}
 ban_msg: ${BAN_MSG}
-whitelist: ${WHITELIST}
+user_whitelist: ${USER_WHITELIST}
+network_whitelist: ${NETWORK_WHITELIST}
 max_unique_streams: ${MAX_UNIQUE_STREAMS}
 
 [telegram]

diff --git a/app/run.py b/app/run.py
@@ -6,6 +6,7 @@
 import json
 import time
 import logging
+import ipaddress
 from pprint import pprint
 from configparser import ConfigParser
 
@@ -109,14 +110,18 @@ def load_bans():
         logging.debug('Loaded bans from disk')
 
 
-def dup_check(user_streams):
+def dup_check(user_streams, network_whitelist):
     """Returns number of unique ip addresses for user"""
     if len(user_streams) == 1:
         return 1
 
     ip_address_list = []
     for stream in user_streams:
-        ip_address_list.append(stream['ip_address'])
+        # only count streams from non-whitelisted ip addresses
+        if any([ ipaddress.IPv4Address(stream['ip_address']) in n for n in network_whitelist ]):
+            logging.debug(f'Ignoring stream from {stream["ip_address"]} (whitelisted)')
+        else:
+            ip_address_list.append(stream['ip_address'])
 
     # return count of unique ip addresses for user
     return len(list(set(ip_address_list)))
@@ -195,7 +200,8 @@ def telegram_notify(message, telegram_bot_key, chat_id):
     max_unique_streams = int(config.get('main', 'max_unique_streams'))
     ban_length_hrs = int(config.get('main', 'ban_length_hrs'))
     ban_msg = config.get('main', 'ban_msg')
-    whitelist = config.get('main', 'whitelist').lower().split()
+    user_whitelist = config.get('main', 'user_whitelist').lower().split()
+    network_whitelist = [ ipaddress.IPv4Network(n) for n in config.get('main', 'network_whitelist').split() ]
     telegram_bot_key = config.get('telegram', 'bot_key')
     telegram_chat_id = config.get('telegram', 'chat_id')
 except FileNotFoundError as err:
@@ -214,7 +220,7 @@ def telegram_notify(message, telegram_bot_key, chat_id):
 
         for user in streams:
             # continue if the user is in a whitelist
-            if user.lower() in whitelist:
+            if user.lower() in user_whitelist:
                 logging.debug(f"User {user} is in whitelist. Not going to count streams")
                 continue
 
@@ -232,7 +238,7 @@ def telegram_notify(message, telegram_bot_key, chat_id):
                     telegram_notify(f"Removed {user} from ban list", telegram_bot_key, telegram_chat_id)
 
             # check to see if user needs to be banned
-            uniq_stream_locations = dup_check(streams[user])
+            uniq_stream_locations = dup_check(streams[user], network_whitelist)
             if uniq_stream_locations > max_unique_streams:
                 logging.info(f"Banning user {user} for {ban_length_hrs} hours for streaming from {uniq_stream_locations} unique locations")
                 ban_list = ban_user(user, ban_length_hrs, ban_list)

diff --git a/docker-compose.yml.EXAMPLE b/docker-compose.yml.EXAMPLE
@@ -11,7 +11,8 @@ services:
       MAX_UNIQUE_STREAMS: 1
       BAN_LENGTH_HRS: 48
       BAN_MSG: YOU HAVE BEEN BANNED FROM PLEX FOR 48 HOURS FOR ACCOUNT SHARING. Please ask @AdminNameHere if you have any questions. This is an automated message.
-      WHITELIST:
+      USER_WHITELIST:
+      NETWORK_WHITELIST:
       PLEX_URL: http://127.0.0.1:10400
       PLEX_TOKEN: ** Plex Token - https://bit.ly/34FeMCo **
       TELEGRAM_BOT_KEY: ** Telegram Bot Key - https://bit.ly/33GhZjV **