Colourized output
[pyderasn.git] / pyderasn.py
index cd96a33200a1a6a7a801bc29b6aa95ca48b860cf..e69b19056ce010fa0ea09e4a111fd68964c34c95 100755 (executable)
@@ -329,7 +329,7 @@ of ``PKIResponse``::
         (
             (
                 "content",
-                decode_path_defby(id_signedData),
+                DecodePathDefBy(id_signedData),
                 "encapContentInfo",
                 "eContentType",
             ),
@@ -341,10 +341,10 @@ of ``PKIResponse``::
         (
             (
                 "content",
-                decode_path_defby(id_signedData),
+                DecodePathDefBy(id_signedData),
                 "encapContentInfo",
                 "eContent",
-                decode_path_defby(id_cct_PKIResponse),
+                DecodePathDefBy(id_cct_PKIResponse),
                 "controlSequence",
                 any,
                 "attrType",
@@ -358,7 +358,7 @@ of ``PKIResponse``::
         ),
     ))
 
-Pay attention for :py:func:`pyderasn.decode_path_defby` and ``any``.
+Pay attention for :py:class:`pyderasn.DecodePathDefBy` and ``any``.
 First function is useful for path construction when some automatic
 decoding is already done. ``any`` means literally any value it meet --
 useful for SEQUENCE/SET OF-s.
@@ -486,6 +486,13 @@ from six import text_type
 from six.moves import xrange as six_xrange
 
 
+try:
+    from termcolor import colored
+except ImportError:
+    def colored(what, *args):
+        return what
+
+
 __all__ = (
     "Any",
     "BitString",
@@ -493,8 +500,8 @@ __all__ = (
     "Boolean",
     "BoundsError",
     "Choice",
-    "decode_path_defby",
     "DecodeError",
+    "DecodePathDefBy",
     "Enumerated",
     "GeneralizedTime",
     "GeneralString",
@@ -1001,10 +1008,24 @@ class Obj(object):
         return self.expl_tlen + self.expl_llen + self.expl_vlen
 
 
-def decode_path_defby(defined_by):
+class DecodePathDefBy(object):
     """DEFINED BY representation inside decode path
     """
-    return "DEFINED BY (%s)" % defined_by
+    __slots__ = ('defined_by',)
+
+    def __init__(self, defined_by):
+        self.defined_by = defined_by
+
+    def __eq__(self, their):
+        if not isinstance(their, self.__class__):
+            return False
+        return self.defined_by == their.defined_by
+
+    def __str__(self):
+        return "DEFINED BY " + str(self.defined_by)
+
+    def __repr__(self):
+        return "<%s: %s>" % (self.__class__.__name__, self.defined_by)
 
 
 ########################################################################
@@ -1072,49 +1093,75 @@ def _pp(
     )
 
 
-def pp_console_row(pp, oids=None, with_offsets=False, with_blob=True):
+def _colorize(what, colour, with_colours, attrs=("bold",)):
+    return colored(what, colour, attrs=attrs) if with_colours else what
+
+
+def pp_console_row(
+        pp,
+        oids=None,
+        with_offsets=False,
+        with_blob=True,
+        with_colours=False,
+):
     cols = []
     if with_offsets:
-        cols.append("%5d%s [%d,%d,%4d]" % (
+        col = "%5d%s" % (
             pp.offset,
             (
                 "  " if pp.expl_offset is None else
                 ("-%d" % (pp.offset - pp.expl_offset))
             ),
-            pp.tlen,
-            pp.llen,
-            pp.vlen,
-        ))
+        )
+        cols.append(_colorize(col, "red", with_colours, ()))
+        col = "[%d,%d,%4d]" % (pp.tlen, pp.llen, pp.vlen)
+        cols.append(_colorize(col, "green", with_colours, ()))
     if len(pp.decode_path) > 0:
         cols.append(" ." * (len(pp.decode_path)))
-        cols.append("%s:" % pp.decode_path[-1])
+        ent = pp.decode_path[-1]
+        if isinstance(ent, DecodePathDefBy):
+            cols.append(_colorize("DEFINED BY", "red", with_colours, ("reverse",)))
+            value = str(ent.defined_by)
+            if (
+                    oids is not None and
+                    ent.defined_by.asn1_type_name ==
+                    ObjectIdentifier.asn1_type_name and
+                    value in oids
+            ):
+                cols.append(_colorize("%s:" % oids[value], "green", with_colours))
+            else:
+                cols.append(_colorize("%s:" % value, "white", with_colours))
+        else:
+            cols.append(_colorize("%s:" % ent, "yellow", with_colours))
     if pp.expl is not None:
         klass, _, num = pp.expl
-        cols.append("[%s%d] EXPLICIT" % (TagClassReprs[klass], num))
+        col = "[%s%d] EXPLICIT" % (TagClassReprs[klass], num)
+        cols.append(_colorize(col, "blue", with_colours))
     if pp.impl is not None:
         klass, _, num = pp.impl
-        cols.append("[%s%d]" % (TagClassReprs[klass], num))
+        col = "[%s%d]" % (TagClassReprs[klass], num)
+        cols.append(_colorize(col, "blue", with_colours))
     if pp.asn1_type_name.replace(" ", "") != pp.obj_name.upper():
-        cols.append(pp.obj_name)
-    cols.append(pp.asn1_type_name)
+        cols.append(_colorize(pp.obj_name, "magenta", with_colours))
+    cols.append(_colorize(pp.asn1_type_name, "cyan", with_colours))
     if pp.value is not None:
         value = pp.value
+        cols.append(_colorize(value, "white", with_colours))
         if (
                 oids is not None and
                 pp.asn1_type_name == ObjectIdentifier.asn1_type_name and
                 value in oids
         ):
-            value = "%s (%s)" % (oids[value], pp.value)
-        cols.append(value)
+            cols.append(_colorize("(%s)" % oids[value], "green", with_colours))
     if with_blob:
         if isinstance(pp.blob, binary_type):
             cols.append(hexenc(pp.blob))
         elif isinstance(pp.blob, tuple):
             cols.append(", ".join(pp.blob))
     if pp.optional:
-        cols.append("OPTIONAL")
+        cols.append(_colorize("OPTIONAL", "red", with_colours))
     if pp.default:
-        cols.append("DEFAULT")
+        cols.append(_colorize("DEFAULT", "red", with_colours))
     return " ".join(cols)
 
 
@@ -1133,7 +1180,7 @@ def pp_console_blob(pp):
         yield " ".join(cols + [", ".join(pp.blob)])
 
 
-def pprint(obj, oids=None, big_blobs=False):
+def pprint(obj, oids=None, big_blobs=False, with_colours=False):
     """Pretty print object
 
     :param Obj obj: object you want to pretty print
@@ -1142,6 +1189,8 @@ def pprint(obj, oids=None, big_blobs=False):
     :param big_blobs: if large binary objects are met (like OctetString
                       values), do we need to print them too, on separate
                       lines
+    :param with_colours: colourize output, if ``termcolor`` library
+                         is available
     """
     def _pprint_pps(pps):
         for pp in pps:
@@ -1152,11 +1201,18 @@ def pprint(obj, oids=None, big_blobs=False):
                         oids=oids,
                         with_offsets=True,
                         with_blob=False,
+                        with_colours=with_colours,
                     )
                     for row in pp_console_blob(pp):
                         yield row
                 else:
-                    yield pp_console_row(pp, oids=oids, with_offsets=True)
+                    yield pp_console_row(
+                        pp,
+                        oids=oids,
+                        with_offsets=True,
+                        with_blob=True,
+                        with_colours=with_colours,
+                    )
             else:
                 for row in _pprint_pps(pp):
                     yield row
@@ -2009,7 +2065,7 @@ class BitString(Obj):
         defined_by, defined = self.defined or (None, None)
         if defined_by is not None:
             yield defined.pps(
-                decode_path=decode_path + (decode_path_defby(defined_by),)
+                decode_path=decode_path + (DecodePathDefBy(defined_by),)
             )
 
 
@@ -2233,7 +2289,7 @@ class OctetString(Obj):
         defined_by, defined = self.defined or (None, None)
         if defined_by is not None:
             yield defined.pps(
-                decode_path=decode_path + (decode_path_defby(defined_by),)
+                decode_path=decode_path + (DecodePathDefBy(defined_by),)
             )
 
 
@@ -3545,7 +3601,7 @@ class Any(Obj):
         defined_by, defined = self.defined or (None, None)
         if defined_by is not None:
             yield defined.pps(
-                decode_path=decode_path + (decode_path_defby(defined_by),)
+                decode_path=decode_path + (DecodePathDefBy(defined_by),)
             )
 
 
@@ -3871,7 +3927,7 @@ class Sequence(Obj):
                     for i, _value in enumerate(value):
                         sub_sub_decode_path = sub_decode_path + (
                             str(i),
-                            decode_path_defby(defined_by),
+                            DecodePathDefBy(defined_by),
                         )
                         defined_value, defined_tail = defined_spec.decode(
                             memoryview(bytes(_value)),
@@ -3899,14 +3955,14 @@ class Sequence(Obj):
                             if value.expled else (value.tlen + value.llen)
                         ),
                         leavemm=True,
-                        decode_path=sub_decode_path + (decode_path_defby(defined_by),),
+                        decode_path=sub_decode_path + (DecodePathDefBy(defined_by),),
                         ctx=ctx,
                     )
                     if len(defined_tail) > 0:
                         raise DecodeError(
                             "remaining data",
                             klass=self.__class__,
-                            decode_path=sub_decode_path + (decode_path_defby(defined_by),),
+                            decode_path=sub_decode_path + (DecodePathDefBy(defined_by),),
                             offset=offset,
                         )
                     value.defined = (defined_by, defined_value)
@@ -4397,7 +4453,7 @@ def generic_decoder():  # pragma: no cover
         __slots__ = ()
         schema = choice
 
-    def pprint_any(obj, oids=None):
+    def pprint_any(obj, oids=None, with_colours=False):
         def _pprint_pps(pps):
             for pp in pps:
                 if hasattr(pp, "_fields"):
@@ -4411,6 +4467,7 @@ def generic_decoder():  # pragma: no cover
                         oids=oids,
                         with_offsets=True,
                         with_blob=False,
+                        with_colours=with_colours,
                     )
                     for row in pp_console_blob(pp):
                         yield row
@@ -4442,6 +4499,11 @@ def main():  # pragma: no cover
         "--defines-by-path",
         help="Python path to decoder's defines_by_path",
     )
+    parser.add_argument(
+        "--with-colours",
+        action='store_true',
+        help="Enable coloured output",
+    )
     parser.add_argument(
         "DERFile",
         type=argparse.FileType("rb"),
@@ -4465,7 +4527,11 @@ def main():  # pragma: no cover
             {"defines_by_path": obj_by_path(args.defines_by_path)}
         ),
     )
-    print(pprinter(obj, oids=oids))
+    print(pprinter(
+        obj,
+        oids=oids,
+        with_colours=True if args.with_colours else False,
+    ))
     if tail != b"":
         print("\nTrailing data: %s" % hexenc(tail))