refactor to support stdin, output options, and better annotation handling

junos_converter.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python3
 
 #
 # Copyright (c) 2017 carles.kishimoto@gmail.com
@@ -16,23 +16,52 @@
 # the License.
 
 import argparse
+import sys
 
 
-def print_set_command(lcommands, leaf):
-    print(("%s %s" % (" ".join(lcommands), leaf)))
+def print_set_command(lcommands, leaf, output_file=None):
+    output = "%s %s" % (" ".join(lcommands), leaf)
+    if output_file:
+        print(output, file=output_file)
+    else:
+        print(output)
+
+
+def print_annotation(annotation_target, annotation_text, output_file=None):
+    # Clean up the annotation text - remove /* and */ and extra whitespace
+    clean_text = annotation_text.strip()
+    if clean_text.startswith('/*'):
+        clean_text = clean_text[2:].strip()
+    if clean_text.endswith('*/'):
+        clean_text = clean_text[:-2].strip()
+
+    # Replace multiple whitespace/newlines with \n for proper formatting
+    import re
+    clean_text = re.sub(r'\s*\*\s*', ' ', clean_text)  # Remove * line prefixes
+    clean_text = re.sub(r'\s+', ' ', clean_text)       # Normalize whitespace
+    clean_text = clean_text.strip()
+
+    output = 'annotate %s "%s"' % (annotation_target, clean_text)
+    if output_file:
+        print(output, file=output_file)
+    else:
+        print(output)
 
 
 def replace_curly(s):
     return s.replace("{", "{\n").replace("}", "\n}")
 
 
-def get_set_config(filein, ignore_annotations):
+def get_set_config(input_source, ignore_annotations, output_file=None):
     try:
-        with open(filein, "r") as f:
-            data = f.read()
+        if input_source == sys.stdin:
+            data = input_source.read()
+        else:
+            with open(input_source, "r") as f:
+                data = f.read()
     except IOError:
-        print("Error: Could not read input file:", filein)
-        exit()
+        print("Error: Could not read input file:", input_source, file=sys.stderr)
+        sys.exit(1)
 
     # Add \n for one-line configs
     if not '"' in data:
@@ -48,56 +77,70 @@ def get_set_config(filein, ignore_annotations):
             ]
         )
 
-    # Keep a list of annotations to be printed at the end
-    lannotations = []
-    annotation = ""
+    # Keep pending annotation for the next element
+    pending_annotation = ""
     lres = ["set"]
+
     for elem in data.split("\n"):
         elem = elem.strip()
         if elem == "" or elem.startswith("#"):
             continue
 
         if elem.startswith("/*"):
-            # Store current annotation
-            annotation = elem.replace("/* ", '"').replace(" */", '"')
+            # Start collecting annotation - may span multiple lines
+            if elem.endswith("*/"):
+                # Single line comment
+                pending_annotation = elem
+            else:
+                # Multi-line comment start
+                pending_annotation = elem
+        elif pending_annotation and not pending_annotation.endswith("*/"):
+            # Continue collecting multi-line comment
+            pending_annotation += " " + elem
         else:
+            # Process regular configuration element
             clean_elem = elem.strip("\t\n\r{ ")
-            if annotation:
-                lannotations.append("top")
+
+            # If we have a pending annotation, apply it to this element
+            if pending_annotation and not ignore_annotations:
+                # Determine the annotation target based on current context
                 if len(lres) > 1:
-                    level = lres[:]
-                    # Replace "set" with "edit"
-                    level[0] = "edit"
-                    lannotations.append("%s" % " ".join(level))
-                # Annotation in a leaf, keep only the keyword
-                clean_elem_orig = clean_elem
-                if ";" in clean_elem:
-                    clean_elem = clean_elem.split()[0]
-                lannotations.append("annotate %s %s" % (clean_elem, annotation))
-                clean_elem = clean_elem_orig
-                annotation = ""
+                    # We're inside a configuration block, annotate relative to parent
+                    annotation_target = " ".join(lres[1:])  # Skip "set"
+                else:
+                    # We're at top level, use the element itself
+                    if ";" in clean_elem:
+                        # It's a leaf, use just the keyword
+                        annotation_target = clean_elem.split()[0]
+                    else:
+                        # It's a container
+                        annotation_target = clean_elem
+
+                print_annotation(annotation_target, pending_annotation, output_file)
+                pending_annotation = ""
+
+            # Handle inactive elements
             if "inactive" in clean_elem:
                 clean_elem = clean_elem.replace("inactive: ", "")
                 linactive = list(lres)
                 linactive[0] = "deactivate"
-                print_set_command(linactive, clean_elem)
+                print_set_command(linactive, clean_elem, output_file)
+
+            # Handle protected elements
             if "protect" in clean_elem:
                 clean_elem = clean_elem.replace("protect: ", "")
                 lprotect = list(lres)
                 lprotect[0] = "protect"
-                print_set_command(lprotect, clean_elem)
+                print_set_command(lprotect, clean_elem, output_file)
+
+            # Handle configuration elements
             if ";" in clean_elem:  # this is a leaf
-                print_set_command(lres, clean_elem.split(";")[0])
+                print_set_command(lres, clean_elem.split(";")[0], output_file)
             elif clean_elem == "}":  # Up one level remove parent
                 lres.pop()
             else:
                 lres.append(clean_elem)
 
-    if not ignore_annotations:
-        # Print all annotations at the end
-        for a in lannotations:
-            print(a)
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description=">>> Juniper display set")
@@ -110,10 +153,37 @@ def get_set_config(filein, ignore_annotations):
     )
     parser.add_argument(
         "--input",
-        required=True,
+        required=False,
+        type=str,
+        help="Specify the input Junos configuration file (if not provided, reads from STDIN)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        required=False,
         type=str,
-        help="Specify the input Junos configuration file",
+        help="Specify the output file (if not provided, writes to STDOUT)",
     )
+
     args = parser.parse_args()
 
-    get_set_config(args.input, args.ignore_annotations)
+    # Determine input source
+    if args.input:
+        input_source = args.input
+    else:
+        input_source = sys.stdin
+
+    # Determine output destination
+    output_file = None
+    if args.output:
+        try:
+            output_file = open(args.output, 'w')
+        except IOError:
+            print("Error: Could not open output file:", args.output, file=sys.stderr)
+            sys.exit(1)
+
+    try:
+        get_set_config(input_source, args.ignore_annotations, output_file)
+    finally:
+        if output_file:
+            output_file.close()