refactoring

wallet.py

@@ -1,7 +1,10 @@
+import gc
 import os
 import json
 import getpass
 from pathlib import Path
+
+import rich
 from web3 import Web3
 from Crypto.Cipher import AES
 from Crypto.Protocol.KDF import PBKDF2
@@ -14,7 +17,6 @@ WALLETS_DIR = "wallets"
 
 web3 = None
 selected_wallet = None
-wallet_data = None
 
 if not os.path.exists(WALLETS_DIR):
     os.makedirs(WALLETS_DIR)
@@ -79,33 +81,32 @@ def save_wallet(account):
 
 
 def select_wallet():
-    global selected_wallet, wallet_data
+    global selected_wallet
     files = os.listdir(WALLETS_DIR)
     if not files:
         print("[red]Нет доступных кошельков.[/]")
         return
 
     for i, fname in enumerate(files):
-        print(f"[{i}] {Path(fname).stem}")
+        print(f"[orange1][{i}] {Path(fname).stem}[/]")
 
     try:
-        choice = int(Prompt.ask("Выберите номер кошелька"))
+        choice = int(Prompt.ask("[bold cyan]Выберите номер кошелька[/]"))
     except:
         print("[red]Неверный номер.[/]")
         return
 
     with open(os.path.join(WALLETS_DIR, files[choice])) as f:
-        wallet_data = json.load(f)
-        selected_wallet = wallet_data['address']
-        print(f"[green]Кошелек выбран: {selected_wallet}[/]")
+        selected_wallet = json.load(f)
+        print(f"[green]Кошелек выбран: {selected_wallet['address']}[/]")
 
 
 def get_balance():
     if not selected_wallet:
         print("[red]Кошелек не выбран.[/]")
         return
-    balance = web3.eth.get_balance(selected_wallet)
-    print(f"Баланс: {web3.from_wei(balance, 'ether')} ETH")
+    balance = web3.eth.get_balance(selected_wallet['address'])
+    print(f"[orange1]Баланс: {web3.from_wei(balance, 'ether')} ETH")
 
 
 def send_transaction():
@@ -113,15 +114,6 @@ def send_transaction():
         print("[red]Кошелек не выбран.[/]")
         return
 
-    print("[bold cyan]Введите пароль от кошелька: [/]", end='')
-    password = getpass.getpass(prompt="")
-
-    try:
-        private_key = decrypt_wallet(wallet_data['data'], password)
-    except:
-        print("[red]Неверный пароль![/]")
-        return
-
     to = Prompt.ask("[bold cyan]Введите адрес получателя[/]")
     value = float(Prompt.ask("[bold cyan]Введите сумму в ETH[/]"))
 
@@ -130,11 +122,11 @@ def send_transaction():
     max_priority_fee = web3.to_wei(priority_fee_gwei, 'gwei')
     max_fee = base_fee + max_priority_fee * 2
 
-    print(f"[yellow]Максимальная комиссия за транзакцию: {float(max_fee) / 10**9} gwei[/]")
+    print(f"[yellow]Максимальная комиссия за транзакцию: {float(max_fee) / 10 ** 9} gwei[/]")
 
     tx = {
         'chainId': web3.eth.chain_id,
-        'nonce': web3.eth.get_transaction_count(selected_wallet),
+        'nonce': web3.eth.get_transaction_count(selected_wallet['address']),
         'to': to,
         'value': web3.to_wei(value, 'ether'),
         'gas': 21000,
@@ -142,13 +134,25 @@ def send_transaction():
         'maxPriorityFeePerGas': max_priority_fee
     }
 
+    print("[bold cyan]Введите пароль от кошелька: [/]", end='')
+    password = getpass.getpass(prompt="")
+
+    try:
+        private_key = decrypt_wallet(selected_wallet['data'], password)
+        del password
+    except:
+        print("[red]Неверный пароль![/]")
+        return
+
     try:
         signed = web3.eth.account.sign_transaction(tx, private_key)
+        del private_key
         tx_hash = web3.eth.send_raw_transaction(signed.raw_transaction)
-
         print(f"[green]Транзакция отправлена! TX Hash: {tx_hash.hex()}[/]")
     except Exception as error:
         print(f"[red]{error.message}[/]")
+    finally:
+        gc.collect()
 
 
 def main():
@@ -161,10 +165,10 @@ def main():
     s. Отправка средств
     i. Импорт существующего кошелька
     q. Выход
-            """)
+    """)
 
     while True:
-        choice = Prompt.ask("Выберите действие").lower()
+        choice = Prompt.ask("[slate_blue1]Выберите действие[/]").lower()
         if choice == 'a':
             connect_to_network()
         elif choice == 'c':