Redo mod-printer with additional whiteout lpstat info

staff/lab/mod-printer

@@ -8,28 +8,9 @@ import sys
 
 def modify_printer(args):
     """Perform internal call to CUPS CLI to modify class(es)."""
-    # Determine which CUPS classes are available
-    lpstat_proc = subprocess.run(['lpstat', '-p'], check=True,
-                                 stdout=subprocess.PIPE, stdin=subprocess.PIPE)
-    lpstat_out = lpstat_proc.stdout.decode('utf-8').strip().split('\n')
-
-    # Parse lpstat output for printer class names
-    # Output is of the form: printer <printername> is idle. <etc>
-    cups_classes = [re.search(r'printer (\S*) is', classname).group(1) for classname in lpstat_out]
-
-    if args.verbose:
-        print('Found CUPS classes: %s' % cups_classes)
-
-    # Override with specified classname if provided and valid
-    if args.classname:
-        if args.classname in cups_classes:
-            cups_classes = [args.classname]
-            if args.verbose:
-                print('Overriding CUPS classes with parameter: %s' % args.classname)
-        else:
-            print('ERROR: Specified class %s could not be found in available classes %s'
-                  % (args.classname, cups_classes))
-            return 1
+    # Get available CUPS classes
+    lpstat_out = do_internal_call(['lpstat', '-c'])
+    cups_classes = parse_cups_classes(lpstat_out, args)
 
     # Perform pre-call logging and input processing
     printer = args.printer.lower()
@@ -40,40 +21,130 @@ def modify_printer(args):
         print('Action to be performed (and flag): %s (%s)' % (args.action, action_flag))
 
     # Call lpadmin command to perform addition/removal
-    for classname in cups_classes:
-        action_cmd = ['lpadmin', '-p', printer + '-' + classname, action_flag, classname]
+    for class_name, class_members in cups_classes.items():
+        printer_name = printer + '-' + class_name
+        if args.action == 'remove' and printer_name not in class_members:
+            # Allow the user to continue - supports removal if only one class was modified
+            print('ERROR: Printer %s is not a member of class %s' % (printer_name, class_name))
+            continue
+        action_cmd = ['lpadmin', '-p', printer_name, action_flag, class_name]
         subprocess.run(action_cmd, check=True)
         action_out = 'added to' if args.action == 'add' else 'removed from'
-        print('Printer %s was successfully %s %s' % (printer, action_out, classname))
+        print('Printer %s was successfully %s %s' % (printer, action_out, class_name))
 
 
 def list_printers(args):
     """Perform internal call to CUPS CLI to list printer status(es) (and potentially jobs)."""
-    # Initialize base lpstat command
-    lpstat_cmd = ['lpstat', '-c']
+    # Get available CUPS classes
+    lpstat_out = do_internal_call(['lpstat', '-c'])
+    cups_classes = parse_cups_classes(lpstat_out, args)
 
-    # Add on to lpstat command with arguments
-    if args.classname:
-        lpstat_cmd.append(args.classname)
-    if args.jobs:
-        lpstat_cmd.append('-o')
+    # Get all CUPS printers
+    lpstat_proc = subprocess.run(['lpstat', '-p'], check=True,
+                                 stdout=subprocess.PIPE, stdin=subprocess.PIPE)
+    lpstat_out = lpstat_proc.stdout.decode('utf-8').strip().split('\n')
+
+    printer = args.printer and args.printer.lower()
+
+    # Parse lpstat output to find all available printers
+    # This includes classes and printers (i.e. shared and non-shared printers)
+    all_printers = []
+    for printer_line in lpstat_out:
+        printer_name_match = re.search(r'printer (\S*)', printer_line)
+        if not printer_name_match:
+            continue
+        printer_name = printer_name_match.group(1)
+
+        printer_arg = (printer and printer in printer_name) or not printer
+        class_arg = (args.classname and args.classname in printer_name) or not args.classname
+
+        if printer_name not in cups_classes and printer_arg and class_arg:
+            all_printers.append(printer_name)
 
     if args.verbose:
-        print('Calling %s to list printer status' % lpstat_cmd)
-
-    # Call lpstat - prints members of the class and can
-    # be configured to include jobs as well. No way to
-    # check for printers not in classes.
-    lpstat_proc = subprocess.run(lpstat_cmd, check=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE)
-
-    # Parse output to fit printer arg, or send directly to stdout
-    status_msgs = lpstat_proc.stdout.decode('utf-8')
-    if args.printer:
-        for line in status_msgs.split('\n'):
-            if args.printer in line or 'members of class' in line:
-                print(line)
+        print('All printers found: %s' % all_printers)
+        print('CUPS class-printer mapping found: %s' % cups_classes)
+
+    # Find differences between all printers and assigned printers
+    diff_printers = all_printers[:]
+    for class_name, class_members in cups_classes.items():
+        for printer_name in class_members[:]:
+            printer_arg = (printer and printer in printer_name) or not printer
+            class_arg = (args.classname and args.classname in printer_name) or not args.classname
+            if printer_arg and class_arg:
+                diff_printers.remove(printer_name)
+            if printer and printer not in printer_name:
+                class_members.remove(printer_name)
+
+        # Print out all printers belonging to class
+        print('Class %s contains printers: %s' % (class_name, ', '.join(class_members)))
+
+    # Print out all printers not assigned to a class (i.e. out of service)
+    if len(diff_printers):
+        print('No classes contain: %s' % ', '.join(diff_printers))
     else:
-        print(status_msgs, end='')
+        print('All printers are in service')
+
+    # Print out jobs currently queued for specified printers (default: all)
+    if args.jobs:
+        jobs_cmd = ['lpstat', '-o']
+        # Specify targets to lpstat command based on classname and printer args
+        if args.classname:
+            jobs_cmd.extend([name for name in all_printers if args.classname in name])
+            jobs_cmd.append(args.classname)
+        if printer:
+            jobs_cmd.extend([name for name in all_printers if printer in name])
+        if args.verbose:
+            print('Printing jobs with command %s' % jobs_cmd)
+
+        # Perform internal jobs call and print result
+        jobs_out = do_internal_call(jobs_cmd)
+        # Remove blank lines/extraneous whitespace
+        jobs_out = [job for job in jobs_out if len(job)]
+        if len(jobs_out):
+            print('\nCurrent print jobs:')
+            print('\n'.join(jobs_out))
+        else:
+            print('\nNo current jobs queued')
+
+
+def do_internal_call(cmd):
+    """Perform a call to the system to retrieve data."""
+    # Perform internal call to get raw data
+    proc = subprocess.run(cmd, check=True,
+                          stdout=subprocess.PIPE, stdin=subprocess.PIPE)
+    # Split data by newline, no trailing or leading whitespace
+    return proc.stdout.decode('utf-8').strip().split('\n')
+
+
+def parse_cups_classes(lpstat_raw, args):
+    """Parse raw CUPS class data into class names and members."""
+    # Data output in dict with class names as keys and class
+    # members as values (each stored in an array)
+    cups_classes = {}
+    curr_classname = None
+    for line in lpstat_raw:
+        match = re.search(r'members of class (\S*):', line)
+        if match:
+            curr_classname = match.group(1)
+            cups_classes[curr_classname] = []
+        else:
+            cups_classes[curr_classname].append(line.strip())
+
+    if args.verbose:
+        print('Found CUPS classes: %s' % cups_classes)
+
+    # Override with specified classname if provided and valid
+    if args.classname:
+        if args.classname in cups_classes:
+            cups_classes = {args.classname: cups_classes[args.classname]}
+            if args.verbose:
+                print('Overriding CUPS classes with parameter: %s' % args.classname)
+        else:
+            print('ERROR: Specified class %s could not be found in available classes %s'
+                  % (args.classname, cups_classes))
+            return 1
+    return cups_classes
 
 
 def main():