Switch to click CLI and colored output

3 files changed, 83 insertions, 54 deletions

diff --git a/Pipfile b/Pipfile
index 375f555..981c2f7 100644
--- a/Pipfile
+++ b/Pipfile
@@ -8,6 +8,8 @@ cryptography = "*"
 pillow = "*"
 pyzbar = "*"
 pymupdf = "*"
+click = "*"
+colorama = "*"
 
 [dev-packages]
 
diff --git a/Pipfile.lock b/Pipfile.lock
index b61de87..590e454 100644
--- a/Pipfile.lock
+++ b/Pipfile.lock
@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "cdf5770078ddbccea03971d385dd543c729159f2da40215c86311cf4293fa047"
+            "sha256": "ae87784ddff12b0dd87073de786b2bfb79d9660c535ea1b487946b0cfcebecc5"
         },
         "pipfile-spec": 6,
         "requires": {},
@@ -56,6 +56,22 @@
             ],
             "version": "==1.14.5"
         },
+        "click": {
+            "hashes": [
+                "sha256:d2b5255c7c6349bc1bd1e59e08cd12acbbd63ce649f2588755783aa94dfb6b1a",
+                "sha256:dacca89f4bfadd5de3d7489b7c8a566eee0d3676333fbb50030263894c38c0dc"
+            ],
+            "index": "pypi",
+            "version": "==7.1.2"
+        },
+        "colorama": {
+            "hashes": [
+                "sha256:5941b2b48a20143d2267e95b1c2a7603ce057ee39fd88e7329b0c292aa16869b",
+                "sha256:9f47eda37229f68eee03b24b9748937c7dc3868f906e8ba69fbcbdd3bc5dc3e2"
+            ],
+            "index": "pypi",
+            "version": "==0.4.4"
+        },
         "cryptography": {
             "hashes": [
                 "sha256:066bc53f052dfeda2f2d7c195cf16fb3e5ff13e1b6b7415b468514b40b381a5b",
diff --git a/verify.py b/verify.py
index 63506d7..97856c4 100644
--- a/verify.py
+++ b/verify.py
@@ -1,5 +1,5 @@
-import argparse
 import base64
+import click
 import fitz
 import json
 
@@ -8,7 +8,7 @@ from pathlib import Path
 
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import padding
-
+from cryptography.exceptions import InvalidSignature
 from PIL import Image
 from pyzbar import pyzbar
 
@@ -23,33 +23,52 @@ def verify(qr_code_bytes):
 
     data = json.loads(payload)
     payload = payload.decode().encode("utf8")
+
+    details = []
+
     if data["ct"] == 1:
         digest = payload
         for i in range(len(data["p"])):
-            print(f"Details of person number {i+1}:")
-            print(f"\tIsraeli ID Number {data['p'][i]['idl']}")
-            print(f"\tID valid by {data['p'][i]['e']}")
-        print(f"Cert Unique ID {data['id']}")
+            details.append(
+                {
+                    "id_num": data["p"][i]["idl"],
+                    "valid_by": data["p"][i]["e"],
+                    "cert_id": data["id"],
+                }
+            )
     elif data["ct"] == 2:
         h = hashes.Hash(hashes.SHA256())
         h.update(payload)
         digest = h.finalize()
-        print(f"Israeli ID Number {data['idl']}")
-        print(f"ID valid by {data['e']}")
-        print(f"Cert Unique ID {data['id']}")
+        details.append(
+            {
+                "id_num": data["idl"],
+                "valid_by": data["e"],
+                "cert_id": data["id"],
+            }
+        )
+
     else:
-        print("Unsupported certificate type")
+        click.secho("Unsupported certificate type", fg="red")
+        return
+
+    for d in details:
+        click.echo(f"\tIsraeli ID Number {d['id_num']}")
+        click.echo(f"\tID valid by {d['valid_by']}")
+        click.echo(f"\tCert Unique ID {d['cert_id']}")
 
     with open(cert("RamzorQRPubKey.pem"), "rb") as f:
         k = serialization.load_pem_public_key(f.read())
-        k.verify(
-            sig,
-            digest,
-            padding.PKCS1v15(),
-            hashes.SHA256(),
-        )
-
-    print("Valid signature!")
+        try:
+            k.verify(
+                sig,
+                digest,
+                padding.PKCS1v15(),
+                hashes.SHA256(),
+            )
+            click.secho("Valid signature!", fg="green", bold=True)
+        except InvalidSignature:
+            click.secho("Invalid signature!", fg="red", bold=True)
 
 
 def read_qr_code(image_path):
@@ -70,41 +89,33 @@ def read_pdf(pdf_path):
                 return read_qr_code(BytesIO(data))
 
 
-def create_arg_parser():
-    parser = argparse.ArgumentParser("Green Pass QR code verifier")
-    group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument(
-        "-i",
-        "--image-path",
-        type=Path,
-        help="Path to an image with the QR code",
-        default=None,
-    )
-    group.add_argument(
-        "-p",
-        "--pdf-path",
-        type=Path,
-        help="Path to PDF file",
-        default=None,
-    )
-    group.add_argument(
-        "-t",
-        "--txt-path",
-        type=Path,
-        help="Path to decoded QR code textual content",
-        default=None,
-    )
-    return parser
+@click.command()
+@click.option("-p", "--pdf-path", type=click.Path(exists=True), help="Path to PDF file")
+@click.option(
+    "-i",
+    "--image-path",
+    type=click.Path(exists=True),
+    help="Path to an image with the QR code",
+)
+@click.option(
+    "-t",
+    "--txt-path",
+    type=click.Path(exists=True),
+    help="Path to decoded QR code textual content",
+)
+def verify_cmd(pdf_path="", image_path="", txt_path=""):
+    if image_path:
+        verify(read_qr_code(image_path))
+    elif pdf_path:
+        verify(read_pdf(pdf_path))
+    elif txt_path:
+        with open(txt_path, "rb") as f:
+            verify(f.read().strip())
+    else:
+        ctx = click.get_current_context()
+        click.echo(ctx.get_help())
+        ctx.exit()
 
 
 if __name__ == "__main__":
-    parser = create_arg_parser()
-    args = parser.parse_args()
-
-    if args.image_path:
-        verify(read_qr_code(args.image_path))
-    elif args.pdf_path:
-        verify(read_pdf(args.pdf_path))
-    elif args.txt_path:
-        with open(args.txt_path, "rb") as f:
-            verify(f.read().strip())
+    verify_cmd()