Tenant Import Feature (#55)

* Add tenant import feature Signed-off-by: jamshale <[email protected]> * Add additonal configs and raw key support Signed-off-by: jamshale <[email protected]> --------- Signed-off-by: jamshale <[email protected]>
openwallet-foundation · Sep 11, 2024 · 8c810df · 8c810df
1 parent c5cf794
commit 8c810df
Show file tree

Hide file tree

Showing 12 changed files with 613 additions and 50 deletions.
diff --git a/askar_tools/README.md b/askar_tools/README.md
@@ -10,7 +10,6 @@
 poetry install
 ```
 
-
 ### Export Wallet:
 
  * Exports a wallet into a file with a readable json format. This can be useful for debugging or for sharing wallet information with others.
@@ -22,25 +21,53 @@ poetry install
     poetry run askar-tools \
     --strategy export \
     --uri postgres://<username>:<password>@<hostname>:<port>/<dbname> \
-    --base-wallet-name <base wallet name> \
-    --base-wallet-key <base wallet key>
+    --wallet-name <base wallet name> \
+    --wallet-key <base wallet key> \
+    --wallet-key-derivation-method <optional> \
+    --export-filename <optional>
     ```
 
-### Multitenant Wallet - Switch from single wallet to multi wallet:
+### Multi-tenant Wallet - Switch from single wallet to multi wallet:
 
 ##### Prerequisites:
     Backup sub-wallet. This operation will delete the sub-wallet when finished. If the wallet is broken for some reason you will not be able to recover it without a backup.
 
  * Converts the profiles in the sub-wallet to individual wallets and databases.
  * After completion, the sub-wallet will be deleted and the deployment should no longer use the `--multitenancy-config '{"wallet_type": "single-wallet-askar"}'` configuration.
 
-- `export` (Output the contents of a wallet to a json file):
+- `mt-convert-to-mw` (Convert from single wallet to multi-wallet multi-tenant agent):
 
     ```
     poetry run askar-tools \ 
     --strategy mt-convert-to-mw  \ 
     --uri postgres://<username>:<password>@<hostname>:<port>/<dbname> \ 
     --wallet-name <base wallet name> \
     --wallet-key <base wallet key> \ 
+    --wallet-key-derivation-method <optional> \
     --multitenant-sub-wallet-name <optional: custom sub wallet name>
+    ```
+
+### Import Wallet:
+
+ * Imports a wallet from a database location into a multi-tenant multi-wallet admin and database location.
+
+- `tenant-import` (Import a wallet into a multi-wallet multi-tenant agent):
+
+    ```
+    poetry run askar-tools \
+    --strategy tenant-import \
+    --uri postgres://<username>:<password>@<hostname>:<port>/<dbname> \
+    --wallet-name <base wallet name> \
+    --wallet-key <base wallet key> \
+    --wallet-key-derivation-method <optional> \
+    --tenant-uri postgres://<username>:<password>@<hostname>:<port>/<dbname> \
+    --tenant-wallet-name <tenant wallet name> \
+    --tenant-wallet-key <tenant wallet key> \
+    --tenant-wallet-key-derivation-method <optional> \
+    --tenant-wallet-type <optional: default is askar> \
+    --tenant-label <optional: default is None> \
+    --tenant-image-url <optional: default is None> \
+    --tenant-webhook-urls <optional: default is None> \
+    --tenant-extra-settings <optional: default is None> \
+    --tenant-dispatch-type <optional: default is None>
     ```
diff --git a/askar_tools/__main__.py b/askar_tools/__main__.py
@@ -11,27 +11,33 @@
 from .multi_wallet_converter import MultiWalletConverter
 from .pg_connection import PgConnection
 from .sqlite_connection import SqliteConnection
+from .tenant_importer import TenantImporter
 
 
 def config():
     """Parse command line arguments."""
     parser = argparse.ArgumentParser("askar-wallet-tools")
+
+    # Strategy
     parser.add_argument(
         "--strategy",
         required=True,
-        choices=["export", "mt-convert-to-mw"],
+        choices=["export", "mt-convert-to-mw", "tenant-import"],
         help=(
             "Specify migration strategy depending on database type, wallet "
             "management mode, and agent type."
         ),
     )
+
+    # Main wallet
     parser.add_argument(
         "--uri",
         required=True,
         help=("Specify URI of database to be migrated."),
     )
     parser.add_argument(
         "--wallet-name",
+        required=True,
         type=str,
         help=(
             "Specify name of wallet to be migrated for DatabasePerWallet "
@@ -40,12 +46,30 @@ def config():
     )
     parser.add_argument(
         "--wallet-key",
+        required=True,
         type=str,
         help=(
             "Specify key corresponding to the given name of the wallet to "
             "be migrated for database per wallet (export) migration strategy."
         ),
     )
+    parser.add_argument(
+        "--wallet-key-derivation-method",
+        type=str,
+        help=("Specify key derivation method for the wallet. Default is 'ARGON2I_MOD'."),
+    )
+
+    # Export
+    parser.add_argument(
+        "--export-filename",
+        type=str,
+        help=(
+            "Specify the filename to export the data to. Default is 'wallet_export.json'."
+        ),
+        default="wallet_export.json",
+    )
+
+    # Multiwallet conversion
     parser.add_argument(
         "--multitenant-sub-wallet-name",
         type=str,
@@ -55,13 +79,75 @@ def config():
         ),
         default="multitenant_sub_wallet",
     )
+
+    # Tenant import
+    parser.add_argument(
+        "--tenant-uri",
+        help=("Specify URI of the tenant database to be imported."),
+    )
+    parser.add_argument(
+        "--tenant-wallet-name",
+        type=str,
+        help=("Specify name of tenant wallet to be imported."),
+    )
+    parser.add_argument(
+        "--tenant-wallet-key",
+        type=str,
+        help=("Specify key corresponding of the tenant wallet to be imported."),
+    )
+    parser.add_argument(
+        "--tenant-wallet-key-derivation-method",
+        type=str,
+        help=(
+            "Specify key derivation method for the tenant wallet. Default is 'ARGON2I_MOD'."
+        ),
+    )
+    parser.add_argument(
+        "--tenant-wallet-type",
+        type=str,
+        help=(
+            """Specify the wallet type of the tenant wallet. Either 'askar' 
+              or 'askar-anoncreds'. Default is 'askar'."""
+        ),
+    )
+    parser.add_argument(
+        "--tenant-label",
+        type=str,
+        help=("Specify the label for the tenant wallet."),
+    )
+    parser.add_argument(
+        "--tenant-image-url",
+        type=str,
+        help=("Specify the image URL for the tenant wallet."),
+    )
+    parser.add_argument(
+        "--tenant-webhook-urls",
+        type=list,
+        help=("Specify the webhook URLs for the tenant wallet."),
+    )
+    parser.add_argument(
+        "--tenant-extra-settings",
+        type=dict,
+        help=("Specify extra settings for the tenant wallet."),
+    )
+    parser.add_argument(
+        "--tenant-dispatch-type",
+        type=str,
+        help=("Specify the dispatch type for the tenant wallet."),
+    )
+
     args, _ = parser.parse_known_args(sys.argv[1:])
 
-    if args.strategy == "export":
-        if not args.wallet_name:
-            raise ValueError("Wallet name required for export strategy")
-        if not args.wallet_key:
-            raise ValueError("Wallet key required for export strategy")
+    if args.strategy == "tenant-import":
+        if (
+            not args.tenant_uri
+            or not args.tenant_wallet_name
+            or not args.tenant_wallet_key
+        ):
+            parser.error(
+                """For tenant-import strategy, tenant-uri, tenant-wallet-name, and 
+                tenant-wallet-key are required."""
+            )
 
     return args
 
@@ -71,7 +157,19 @@ async def main(
     uri: str,
     wallet_name: Optional[str] = None,
     wallet_key: Optional[str] = None,
+    wallet_key_derivation_method: Optional[str] = "ARGON2I_MOD",
     multitenant_sub_wallet_name: Optional[str] = "multitenant_sub_wallet",
+    tenant_uri: Optional[str] = None,
+    tenant_wallet_name: Optional[str] = None,
+    tenant_wallet_key: Optional[str] = None,
+    tenant_wallet_type: Optional[str] = "askar",
+    tenant_wallet_key_derivation_method: Optional[str] = "ARGON2I_MOD",
+    tenant_label: Optional[str] = None,
+    tenant_image_url: Optional[str] = None,
+    tenant_webhook_urls: Optional[list] = None,
+    tenant_extra_settings: Optional[dict] = None,
+    tenant_dispatch_type: Optional[str] = "default",
+    export_filename: Optional[str] = "wallet_export.json",
 ):
     """Run the main function."""
     logging.basicConfig(level=logging.WARN)
@@ -88,16 +186,49 @@ async def main(
     # Strategy setup
     if strategy == "export":
         await conn.connect()
-        print("wallet_name", wallet_name)
-        method = Exporter(conn=conn, wallet_name=wallet_name, wallet_key=wallet_key)
+        method = Exporter(
+            conn=conn,
+            wallet_name=wallet_name,
+            wallet_key=wallet_key,
+            wallet_key_derivation_method=wallet_key_derivation_method,
+            export_filename=export_filename,
+        )
     elif strategy == "mt-convert-to-mw":
         await conn.connect()
         method = MultiWalletConverter(
             conn=conn,
             wallet_name=wallet_name,
             wallet_key=wallet_key,
+            wallet_key_derivation_method=wallet_key_derivation_method,
             sub_wallet_name=multitenant_sub_wallet_name,
         )
+    elif strategy == "tenant-import":
+        tenant_parsed = urlparse(tenant_uri)
+        if tenant_parsed.scheme == "sqlite":
+            tenant_conn = SqliteConnection(tenant_uri)
+        elif tenant_parsed.scheme == "postgres":
+            tenant_conn = PgConnection(tenant_uri)
+        else:
+            raise ValueError("Unexpected tenant DB URI scheme")
+
+        await conn.connect()
+        await tenant_conn.connect()
+        method = TenantImporter(
+            admin_conn=conn,
+            admin_wallet_name=wallet_name,
+            admin_wallet_key=wallet_key,
+            admin_wallet_key_derivation_method=wallet_key_derivation_method,
+            tenant_conn=tenant_conn,
+            tenant_wallet_name=tenant_wallet_name,
+            tenant_wallet_key=tenant_wallet_key,
+            tenant_wallet_type=tenant_wallet_type,
+            tenant_wallet_key_derivation_method=tenant_wallet_key_derivation_method,
+            tenant_label=tenant_label,
+            tenant_image_url=tenant_image_url,
+            tenant_webhook_urls=tenant_webhook_urls,
+            tenant_extra_settings=tenant_extra_settings,
+            tenant_dispatch_type=tenant_dispatch_type,
+        )
     else:
         raise Exception("Invalid strategy")
 

diff --git a/askar_tools/exporter.py b/askar_tools/exporter.py
@@ -5,6 +5,7 @@
 
 from aries_askar import Store
 
+from .key_methods import KEY_METHODS
 from .pg_connection import PgConnection
 from .sqlite_connection import SqliteConnection
 
@@ -17,17 +18,23 @@ def __init__(
         conn: SqliteConnection | PgConnection,
         wallet_name: str,
         wallet_key: str,
+        wallet_key_derivation_method: str = "ARGON2I_MOD",
+        export_filename: str = "wallet_export.json",
     ):
         """Initialize the Exporter object.
 
         Args:
             conn: The connection object.
             wallet_name: The name of the wallet.
             wallet_key: The key for the wallet.
+            wallet_key_derivation_method: The key derivation method for the wallet.
+            export_filename: The name of the export file.
         """
         self.conn = conn
         self.wallet_name = wallet_name
         self.wallet_key = wallet_key
+        self.wallet_key_derivation_method = wallet_key_derivation_method
+        self.export_filename = export_filename
 
     async def _get_decoded_items_and_tags(self, store):
         scan = store.scan()
@@ -51,18 +58,22 @@ async def _get_decoded_items_and_tags(self, store):
 
     async def export(self):
         """Export the wallet data."""
-        print("Exporting wallet to wallet_export.json...")
+        print(f"Exporting wallet to {self.export_filename}...")
 
         tables = {"config": {}, "items": {}, "profiles": {}}
-        store = await Store.open(self.conn.uri, pass_key=self.wallet_key)
+        store = await Store.open(
+            self.conn.uri,
+            pass_key=self.wallet_key,
+            key_method=KEY_METHODS[self.wallet_key_derivation_method],
+        )
 
         tables["items"] = await self._get_decoded_items_and_tags(store)
 
         tables["config"] = await self.conn.get_root_config()
 
         tables["profiles"] = await self.conn.get_profiles()
 
-        with open("wallet_export.json", "w") as json_file:
+        with open(self.export_filename, "w") as json_file:
             json.dump(tables, json_file, indent=4)
 
         await store.close()

diff --git a/askar_tools/key_methods.py b/askar_tools/key_methods.py
@@ -0,0 +1,7 @@
+""".Key methods for Askar wallet."""
+
+KEY_METHODS = {
+    "RAW": "RAW",
+    "ARGON2I_INT": "kdf:argon2i:int",
+    "ARGON2I_MOD": "kdf:argon2i:mod",
+}
diff --git a/askar_tools/multi_wallet_converter.py b/askar_tools/multi_wallet_converter.py
@@ -3,15 +3,10 @@
 from aries_askar import Store
 
 from .error import ConversionError
+from .key_methods import KEY_METHODS
 from .pg_connection import PgConnection
 from .sqlite_connection import SqliteConnection
 
-KEY_METHODS = {
-    "KEY_DERIVATION_RAW": "RAW",
-    "KEY_DERIVATION_ARGON2I_INT": "kdf:argon2i:int",
-    "KEY_DERIVATION_ARGON2I_MOD": "kdf:argon2i:mod",
-}
-
 
 class MultiWalletConverter:
     """Util class for converting multi-tenant wallets between single wallet and multi wallet."""  # noqa: E501
@@ -21,6 +16,7 @@ def __init__(
         conn: SqliteConnection | PgConnection,
         wallet_name: str,
         wallet_key: str,
+        wallet_key_derivation_method: str,
         sub_wallet_name: str,
     ):
         """Initialize the MultiWalletConverter instance.
@@ -29,11 +25,13 @@ def __init__(
             conn (SqliteConnection): The SQLite connection object.
             wallet_name (str): The name of the wallet.
             wallet_key (str): The key for the wallet.
+            wallet_key_derivation_method (str): The key derivation method for the wallet.
             sub_wallet_name (str): The name of the sub wallet.
         """
         self.conn = conn
         self.admin_wallet_name = wallet_name
         self.admin_wallet_key = wallet_key
+        self.wallet_key_derivation_method = wallet_key_derivation_method
         self.sub_wallet_name = sub_wallet_name
 
     def get_wallet_records(self, entries):
@@ -87,7 +85,7 @@ async def convert_single_wallet_to_multi_wallet(self):
                 )
                 key_method = KEY_METHODS.get(
                     wallet_record["settings"].get(
-                        "wallet.key_derivation_method", "KEY_DERIVATION_ARGON2I_MOD"
+                        "wallet.key_derivation_method", "ARGON2I_MOD"
                     )
                 )
                 print(