native room filter

megolm_backup.py

@@ -20,6 +20,7 @@ import base64
 import getpass
 import hashlib
 import hmac
+import json
 import struct
 import sys
 
@@ -140,26 +141,17 @@ def dec_session_data(passphrase, session_data):
     return cipher.decrypt(ciphertext)
 
 
-def main(args):
+def main():
     parser = argparse.ArgumentParser(description="Operate on megolm session backups.")
-    parser.add_argument("file", nargs="?", default="-", help="Input text file (- for stdin).")
-    parser.add_argument("-o", "--output", default="-", required=False, help="Output text file (- for stdout).")
-    mode_group = parser.add_mutually_exclusive_group(required=True)
-    mode_group.add_argument("--into", dest="mode", const="encrypt", action="store_const",
-                            help="Encrypt and represent file as a megolm session backup.")
-    mode_group.add_argument("--from", dest="mode", const="decrypt", action="store_const",
-                            help="Decrypt the given megolm session and output the contents.")
-    args = parser.parse_args(args)
+    parser.add_argument("file", nargs='?', help="megolm session data")
+    parser.add_argument("room_id", nargs='?', help="Room id to filter (optional)")
+    parser.add_argument("-o", "--output", help="Output to file")
+    parser.add_argument("-p", "--plain", dest="mode", const="plain", action="store_const",
+                        help="Returns the plain unencrypted content")
+    args = parser.parse_args()
 
-    if args.file == "-":
-        args.file = "/dev/stdin"
-    if args.output == "-":
-        args.output = "/dev/stdout"
-
-    action = {
-        "encrypt": enc_session_data,
-        "decrypt": dec_session_data,
-    }[args.mode]
+    if not args.mode:
+        args.mode = "encrypted"
 
     with open(args.file, "rb") as f:
         data = f.read()
@@ -167,13 +159,30 @@ def main(args):
     # Wait until after reading input to get the passphrase so pipelines work
     # properly. This results in slightly strange behaviour for interactive
     # uses, but most people will be using this in a pipeline.
-    passphrase = getpass.getpass("Backup passphrase [mode=%s]: " % (args.mode,))
-    output = action(passphrase, data)
+    passphrase = getpass.getpass(f"Backup passphrase [mode={args.mode}]: ")
+    decrypted_data = dec_session_data(passphrase, data)
 
-    with open(args.output, "wb") as f:
-        f.write(output + b"\n")
-        f.flush()
+    if args.room_id:
+        json_data = json.loads(decrypted_data)
+        filtered_data = [ key for key in json_data if key['room_id'] == args.room_id ]
+
+        if not filtered_data:
+            bail(f"No keys found for room {args.room_id}")
+
+        decrypted_data = str(filtered_data).encode()
+
+    output = decrypted_data
+
+    if args.mode == "encrypted":
+        output = enc_session_data(passphrase, output)
+
+    if args.output:
+        with open(args.output, "wb") as file:
+            file.write(output + b"\n")
+            file.flush()
+    else:
+        print(f"\x1b[1A{output.decode()}")
 
 
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    main()