pyblock/pybitblock/ai/ui.py
GaltRanch 893aabc85d Security audit round 2: eliminate shell=True, mask secrets, fix race conditions
- Replace all shell=True subprocess calls with Python-native processing
  (nodeconnection.py, SPV/nodeconnection.py, SPV/ppi.py)
- Mask sensitive inputs (private keys, passwords, tokens) with getpass
- Add threading.Lock to block_explorer.py shared state
- Use json.loads() instead of fragile string splitting in apisnd.py
- Add path validation before file open in apisnd.py
- Replace random.randint with secrets.randbelow for mining nonces
- Fix destructive exception handlers in clone.py and feed.py
- Replace bare except clauses with specific exceptions + logging
- Remove unused imports (psutil, xmltodict, block_visualizer, base64, say)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-03 11:15:59 -03:00

315 lines
9.5 KiB
Python

"""Terminal UI for PyBLOCK AI Assistant."""
import getpass
import logging
import sys
import time
import qrcode
import requests
from rich.console import Console
from rich.markdown import Markdown
logger = logging.getLogger(__name__)
from shared.display import clear
from pblogo import blogo
from .client import AstrolexisClient
from .context import gather_node_context
_console = Console()
# Colors
G = "\033[1;32;40m" # green
C = "\033[1;36;40m" # cyan
Y = "\033[1;33;40m" # yellow
R = "\033[1;31;40m" # red
W = "\033[1;37;40m" # white bold
D = "\033[0;37;40m" # dim/default
DIM = "\033[2m"
def ai_menu(path, lndconnectload=None):
"""Main AI assistant menu. Requires Astrolexis token in config."""
from config import cfg
token = cfg.settings.get("astrolexis_token", "")
if not token:
token = _setup_token(cfg)
if not token:
return
client = AstrolexisClient(token)
try:
info = client.verify()
except Exception as e:
clear()
blogo()
print(f"\n {R}Error connecting to Astrolexis:{D} {e}")
print(f" Check your token in Settings.\n")
input(" Press Enter to return...")
return
_chat_loop(client, path, lndconnectload, info["balance_sats"])
def _setup_token(cfg):
"""First-time token setup."""
clear()
blogo()
print(f"""
{W}AI Assistant Setup{D}
Powered by {C}Astrolexis KCode{D}
To use the AI Assistant, you need an Astrolexis token.
Get yours at: {Y}https://astrolexis.space/pyblock{D}
Enter your token below, or press Enter to cancel.
""")
token = getpass.getpass(" Token: ").strip()
if not token:
return None
if not token.startswith("astrolexis_"):
print(f"\n {R}Invalid token format.{D} Must start with 'astrolexis_'")
input(" Press Enter to return...")
return None
settings = cfg.settings
settings["astrolexis_token"] = token
cfg.save("pyblocksettings.conf", settings)
print(f"\n {G}Token saved.{D}")
time.sleep(1)
return token
def _render_response(text):
"""Render AI response using Rich Markdown with forced UTF-8 output."""
import io
width = min(80, _console.width - 4)
buf = io.StringIO()
temp = Console(file=buf, width=width, force_terminal=True)
temp.print()
temp.print(Markdown(text), width=width)
temp.print()
rendered = buf.getvalue()
# Write as UTF-8 bytes directly to avoid encoding issues
sys.stdout.buffer.write(rendered.encode('utf-8'))
sys.stdout.buffer.flush()
def _status_line(balance):
"""Compact status line."""
return (
f" {C}AI Assistant{D} | "
f"Balance: {G}{balance:,}{D} sats | "
f"{DIM}T{D}=topup {DIM}U{D}=usage {DIM}C{D}=clear {DIM}Q{D}=quit"
)
def _chat_loop(client, path, lndconnectload, balance):
"""Continuous chat loop — no screen clearing between messages."""
conversation = []
context = None
clear()
blogo()
print(f"""
{W}AI Assistant{D}
Powered by {C}Astrolexis KCode{D}
Balance: {G}{balance:,}{D} sats
{DIM}Ask anything about your Bitcoin/Lightning node.
Commands: T=topup U=usage C=clear Q=quit{D}
""")
# Gather context once at start, refresh on new blocks
try:
context = gather_node_context(path, lndconnectload)
except (requests.RequestException, OSError, ValueError, KeyError) as e:
logger.debug("Initial node context gather failed: %s", e)
context = {}
while True:
try:
# Prompt — distinct color from AI response
user_input = input(f"\n {Y}pyblock>{D} ").strip()
if not user_input:
continue
upper = user_input.upper()
if upper == "Q":
break
if upper == "T":
balance = _topup_flow(client)
# Redraw header after topup
clear()
blogo()
print(f"\n{_status_line(balance)}\n")
continue
if upper == "U":
_show_usage(client)
try:
balance = client.get_balance()
except (requests.RequestException, KeyError, ValueError) as e:
logger.debug("Balance refresh failed: %s", e)
print(f"\n{_status_line(balance)}\n")
continue
if upper == "C":
conversation = []
clear()
blogo()
print(f"\n {DIM}Conversation cleared.{D}\n")
print(f"{_status_line(balance)}\n")
continue
# Add to conversation
conversation.append({"role": "user", "content": user_input})
# Refresh context periodically
try:
context = gather_node_context(path, lndconnectload)
except (requests.RequestException, OSError, ValueError, KeyError) as e:
logger.debug("Node context refresh failed: %s", e)
# Visual separator between user input and AI response
print(f"\n {C}{'' * 60}{D}")
# Stream response
full_response = ""
try:
for chunk in client.chat(
conversation, node_context=context
):
if chunk.get("type") == "content_block_delta":
text = chunk.get("delta", {}).get("text", "")
full_response += text
_render_response(full_response)
print(f" {C}{'' * 60}{D}")
# Add to conversation history
conversation.append({
"role": "assistant", "content": full_response
})
# Update balance
try:
balance = client.get_balance()
except (requests.RequestException, KeyError, ValueError) as e:
logger.debug("Post-chat balance refresh failed: %s", e)
# Show balance below separator
print(f" {DIM}Balance: {balance:,} sats{D}")
except requests.exceptions.HTTPError as e:
if e.response is not None and e.response.status_code == 402:
data = e.response.json()
bal = data.get('balance_sats', 0)
cost = data.get('estimated_cost', '?')
print(
f" {R}Insufficient balance{D} "
f"({bal} sats, need ~{cost})."
)
print(f" Press {Y}T{D} to top up.\n")
conversation.pop()
else:
print(f" {R}Error:{D} {e}\n")
conversation.pop()
except Exception as e:
print(f" {R}Error:{D} {e}\n")
if conversation and conversation[-1]["role"] == "user":
conversation.pop()
except KeyboardInterrupt:
print(f"\n\n {DIM}Ctrl+C — back to main menu{D}\n")
break
except EOFError:
break
def _topup_flow(client):
"""Lightning top-up flow. Returns new balance."""
clear()
blogo()
print(f"""
{W}Top Up Balance{D}
Enter amount in sats (100 - 100,000):
""")
try:
amount = int(input(" Amount: ").strip())
if amount < 100 or amount > 100000:
print(f" {R}Amount must be between 100 and 100,000 sats.{D}")
input(" Press Enter to return...")
return client.get_balance()
except (ValueError, KeyboardInterrupt):
return client.get_balance()
try:
result = client.topup(amount)
except Exception as e:
print(f"\n {R}Error creating invoice:{D} {e}")
input(" Press Enter to return...")
return client.get_balance()
invoice = result["invoice"]
payment_hash = result["payment_hash"]
clear()
blogo()
print(f"\n {W}Lightning Invoice ({amount:,} sats){D}\n")
# QR code
try:
qr = qrcode.QRCode(box_size=1, border=1)
qr.add_data(invoice.upper())
print("\033[1;30;47m")
qr.print_ascii()
print(D)
except (ValueError, OSError) as e:
logger.debug("QR code generation failed: %s", e)
print(f" {invoice}\n")
print(f" Pay with any Lightning wallet. Waiting for payment...\n")
# Poll for payment
for _ in range(200): # ~10 min max
time.sleep(3)
try:
if client.check_payment(payment_hash):
new_balance = client.get_balance()
print(
f"\n {G}Payment received! "
f"New balance: {new_balance:,} sats{D}\n"
)
time.sleep(2)
return new_balance
except (requests.RequestException, KeyError, ValueError) as e:
logger.debug("Payment check failed: %s", e)
sys.stdout.write(".")
sys.stdout.flush()
print(f"\n\n {R}Invoice expired.{D} Try again.")
time.sleep(2)
return client.get_balance()
def _show_usage(client):
"""Display usage statistics inline."""
try:
stats = client.usage(30)
print(f"""
{W}Usage (last 30 days){D}
Queries: {stats.get('total_queries', 0)}
Sats spent: {stats.get('total_sats', 0):,}
Tokens in: {stats.get('total_tokens_in', 0):,}
Tokens out: {stats.get('total_tokens_out', 0):,}
Balance: {G}{stats.get('balance_sats', 0):,}{D} sats
""")
except Exception as e:
print(f"\n {R}Error:{D} {e}\n")