allow username exclusion

Author: jiehanzheng

README.md

@@ -30,6 +30,7 @@ Usage
 ```
 usage: squid2radius.py [-h] [-p RADIUS_ACCT_PORT]
                        [--radius-nasid RADIUS_NASID] [--squid-path SQUID_PATH]
+                       [--exclude-pattern EXCLUDE_PATTERN]
                        logfile_path radius_server radius_secret
 ```
 
@@ -43,8 +44,13 @@ It is certainly a good idea to make a cron job for this.
 
 You should also read [SquidFaq/SquidLogs](http://wiki.squid-cache.org/SquidFaq/SquidLogs#access.log) to make sure your log files are in reasonable sizes.
 
+--exclude-pattern
+-----------------
+
+If for some reason you need to prevent usage information of certain user from being sent to the RADIUS server, there is an argument for that!  Use `--exclude-pattern="(girl|boy)friend"` and squid2radius won't send usage of either your `girlfriend` or `boyfriend` to the RADIUS server.
+
 Note
-----
+====
 
 The script assumes that you are using the default [Squid native access.log format](http://wiki.squid-cache.org/Features/LogFormat#squid) on first ten columns of your log file.  If you need custom columns, add them after the default ones.
 

squid2radius.py

@@ -3,6 +3,7 @@
 import sys
 import argparse
 import time
+import re
 from subprocess import call
 import pyrad.packet
 from pyrad.client import Client
@@ -17,6 +18,10 @@
 parser.add_argument('-p', '--radius-acct-port', default='1813')
 parser.add_argument('--radius-nasid', default='squid')
 parser.add_argument('--squid-path', default='/usr/sbin/squid')
+parser.add_argument('--exclude-pattern', default='', help='do not send to ' \
+                                                          'server if ' \
+                                                          'username contains '\
+                                                          'this regexp')
 args = parser.parse_args()
 
 
@@ -26,7 +31,7 @@
 sys.stdout.write("Analyzing")
 sum_bytes = {}
 for i, line in enumerate(logfile):
-  if i % 1000 == 0: sys.stdout.write('.')
+  if i % 1000 == 0: sys.stdout.write('.'); sys.stdout.flush()
 
   # http://wiki.squid-cache.org/Features/LogFormat
   _, _, _, _, num_bytes, _, _, rfc931, _, _ = line.split()[:10]
@@ -44,10 +49,21 @@
              dict=Dictionary(sys.path[0] + "/dictionary"))
 
 
+if args.exclude_pattern:
+  print "Exclusion check has been enabled."
+  exclude_pattern = re.compile(args.exclude_pattern)
+
+
 print "Sending..."
 for username, total_bytes in sum_bytes.iteritems():
   sys.stdout.write(username + ' ' + str(total_bytes))
   sys.stdout.write('.')
+  sys.stdout.flush()
+  
+  if args.exclude_pattern and exclude_pattern.match(username):
+    sys.stdout.write("..skipped!\n")
+    sys.stdout.flush()
+    continue
 
   session_id = str(time.time())
 
@@ -62,6 +78,7 @@
     raise Exception("mysterious RADIUS server response to Start packet")
 
   sys.stdout.write('.')
+  sys.stdout.flush()
 
   req = srv.CreateAcctPacket()
   req['User-Name'] = username
@@ -75,6 +92,7 @@
     raise Exception("mysterious RADIUS server response to Stop packet")
 
   sys.stdout.write(".\n")
+  sys.stdout.flush()
 
 
 print "\nRotating squid log..."