Rewriting tlnb command to handle naming servers.

utils/tlnb.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python
 
+import argparse
+import json
 import os
 import re
-import time
 import select
 import subprocess
 import sys
-import json
+import time
+
 from requests.exceptions import ConnectionError
 
 sys.path.insert(0, '/srv/newsblur')
@@ -20,7 +22,7 @@ IGNORE_HOSTS = [
 # Use this to count the number of times each user shows up in the logs. Good for finding abusive accounts.
 # tail -n20000 logs/newsblur.log | sed 's/\x1b\[[0-9;]*m//g' | sed -En 's/.*?[0-9]s\] \[([a-zA-Z0-9]+\*?)\].*/\1/p' | sort | uniq -c | sort
 
-def main(roles=None, command=None, path=None):
+def main(hostnames=None, roles=None, command=None, path=None):
     delay = 1
 
     hosts = subprocess.check_output(['ansible-inventory', '--list'])
@@ -29,6 +31,19 @@ def main(roles=None, command=None, path=None):
         return
     hosts = json.loads(hosts)
 
+    if not path:
+        path = "/srv/newsblur/logs/newsblur.log"
+    if not command:
+        command = "tail -f"
+    
+    if hostnames in ['app', 'task', 'push']:
+        roles = hostnames
+        hostnames = None
+        
+    if hostnames:
+        roles = hosts
+        hostnames = validate_hostnames(hostnames.split(','), hosts)
+        
     if not roles:
         roles = ['app']
     if not isinstance(roles, list):
@@ -36,8 +51,15 @@ def main(roles=None, command=None, path=None):
 
     while True:
         try:
-            streams = create_streams_for_roles(hosts, roles, command=command, path=path)
-            print(" --- Loading %s %s Log Tails ---" % (len(streams), roles))
+            if hostnames:
+                streams = []
+                found = set()
+                for host in hostnames:
+                    follow_host(roles[0], streams, found, host, command, path)
+                print(" --- Loading %s %s Log Tails ---" % (len(streams), hostnames))
+            else:
+                streams = create_streams_for_roles(hosts, roles, command=command, path=path)
+                print(" --- Loading %s %s Log Tails ---" % (len(streams), roles))
             read_streams(streams)
         # except UnicodeDecodeError: # unexpected end of data
         #     print " --- Lost connections - Retrying... ---"
@@ -51,15 +73,20 @@ def main(roles=None, command=None, path=None):
         except KeyboardInterrupt:
             print(" --- End of Logging ---")
             break
+        
+def validate_hostnames(hostnames, hosts):
+    validated_hostnames = []
+    for hostname in hostnames:
+        if hostname in hosts['_meta']['hostvars']:
+            validated_hostnames.append(hostname)
+        else:
+            print(f"Hostname {hostname} not found in inventory.")
+    print(f"Validated hostnames: {validated_hostnames}")
+    return validated_hostnames
 
 def create_streams_for_roles(hosts, roles, command=None, path=None):
     streams = list()
     found = set()
-    print(path)
-    if not path:
-        path = "/srv/newsblur/logs/newsblur.log"
-    if not command:
-        command = "tail -f"
     for role in roles:
         if role in hosts:
             for hostname in hosts[role]['hosts']:
@@ -111,4 +138,10 @@ def read_streams(streams):
                 break
 
 if __name__ == "__main__":
-    main(*sys.argv[1:])
+    parser = argparse.ArgumentParser(description='Tail logs from multiple hosts.')
+    parser.add_argument('hostnames', help='Comma-separated list of hostnames', nargs='?')
+    parser.add_argument('roles', help='Comma-separated list of roles', nargs='?')
+    parser.add_argument('--command', help='Command to run on the remote host')
+    parser.add_argument('--path', help='Path to the log file')
+    args = parser.parse_args()
+    main(args.hostnames, command=args.command, path=args.path)