From 04b1224d6d4daa909db23459825c81f316204992 Mon Sep 17 00:00:00 2001
From: cyc60 <avsysoev60@gmail.com>
Date: Mon, 29 Jul 2024 19:04:29 +0300
Subject: [PATCH] Generate keys in tmp directory

Signed-off-by: cyc60 <avsysoev60@gmail.com>
---
 src/commands/create_keys.py | 64 ++++++++++++++++++++++---------------
 src/common/vault_config.py  | 11 +++++++
 2 files changed, 50 insertions(+), 25 deletions(-)

diff --git a/src/commands/create_keys.py b/src/commands/create_keys.py
index 4371e55b..5566e3d0 100644
--- a/src/commands/create_keys.py
+++ b/src/commands/create_keys.py
@@ -62,34 +62,48 @@ def create_keys(
     per_keystore_password: bool,
     pool_size: int | None,
 ) -> None:
-    config = VaultConfig(vault, Path(data_dir))
-    config.load(mnemonic)
+    vault_config = VaultConfig(vault, Path(data_dir))
+    vault_config.load(mnemonic)
 
-    deposit_data_file = config.vault_dir / 'deposit_data.json'
-    keystores_dir = config.vault_dir / 'keystores'
+    deposit_data_file = vault_config.vault_dir / 'deposit_data.json'
+    keystores_dir = vault_config.vault_dir / 'keystores'
     password_file = keystores_dir / 'password.txt'
 
-    credentials = CredentialManager.generate_credentials(
-        network=config.network,
-        vault=vault,
-        mnemonic=mnemonic,
-        count=count,
-        start_index=config.mnemonic_next_index,
-        pool_size=pool_size,
-    )
-    deposit_data = _export_deposit_data_json(
-        credentials=credentials, filename=str(deposit_data_file), pool_size=pool_size
-    )
-
-    _export_keystores(
-        credentials=credentials,
-        keystores_dir=keystores_dir,
-        password_file=password_file,
-        per_keystore_password=per_keystore_password,
-        pool_size=pool_size,
-    )
-
-    config.increment_mnemonic_index(count)
+    # first generate files in tmp directory
+    vault_config.create_tmp_dir()
+    deposit_data_tmp_file = vault_config.vault_tmp_dir / 'deposit_data.json'
+    keystores_tmp_dir = vault_config.vault_tmp_dir / 'keystores'
+
+    try:
+        credentials = CredentialManager.generate_credentials(
+            network=vault_config.network,
+            vault=vault,
+            mnemonic=mnemonic,
+            count=count,
+            start_index=vault_config.mnemonic_next_index,
+            pool_size=pool_size,
+        )
+        deposit_data = _export_deposit_data_json(
+            credentials=credentials, filename=str(deposit_data_tmp_file), pool_size=pool_size
+        )
+
+        _export_keystores(
+            credentials=credentials,
+            keystores_dir=keystores_tmp_dir,
+            password_file=password_file,
+            per_keystore_password=per_keystore_password,
+            pool_size=pool_size,
+        )
+
+        vault_config.increment_mnemonic_index(count)
+
+        # move files from tmp dir
+        deposit_data_tmp_file.replace(deposit_data_file)
+        for src_file in keystores_tmp_dir.glob('*.json'):
+            src_file.rename(keystores_dir.joinpath(src_file.name))
+
+    finally:
+        vault_config.remove_tmp_dir()
 
     click.echo(
         f'Done. Generated {greenify(count)} keys for {greenify(vault)} vault.\n'
diff --git a/src/common/vault_config.py b/src/common/vault_config.py
index 03cb0393..42d84710 100644
--- a/src/common/vault_config.py
+++ b/src/common/vault_config.py
@@ -1,5 +1,6 @@
 import json
 import re
+import shutil
 from pathlib import Path
 
 import click
@@ -28,6 +29,16 @@ def __init__(
     def exists(self) -> bool:
         return self.config_path.is_file()
 
+    @property
+    def vault_tmp_dir(self) -> Path:
+        return self.vault_dir / '.tmp'
+
+    def create_tmp_dir(self) -> None:
+        self.vault_tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    def remove_tmp_dir(self) -> None:
+        shutil.rmtree(self.vault_tmp_dir)
+
     def load(self, mnemonic: str | None = None) -> None:
         if self.config_path.is_file():
             with self.config_path.open('r') as f: