diff --git a/.coverage b/.coverage deleted file mode 100644 index a714e5b..0000000 Binary files a/.coverage and /dev/null differ diff --git a/.gitignore b/.gitignore index 7a60b85..b07531c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ __pycache__/ *.pyc +bot_activity.log +.coverage diff --git a/src/handlers/common.py b/src/handlers/common.py index e5eac65..3aee930 100644 --- a/src/handlers/common.py +++ b/src/handlers/common.py @@ -1,7 +1,8 @@ +import os import logging import functools from telegram import Update -from telegram.ext import CallbackContext +from telegram.ext import CallbackContext, ConversationHandler from config import TIMEOUT_NAME, TIMEOUT_FIRE_NAME @@ -22,6 +23,39 @@ async def wrapper(update: Update, context: CallbackContext, *args, **kwargs): return wrapper +def cb_auth_required(func): + """Decorator to restrict callback query handler access to allowed users only. + Reads the whitelist from context.bot_data['allowed_user_ids']. + Returns ConversationHandler.END when access is denied so the conversation + is cleanly terminated regardless of the current state. + """ + @functools.wraps(func) + async def wrapper(update: Update, context: CallbackContext, *args, **kwargs): + query = update.callback_query + user_id = str(query.from_user.id) + allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) + if user_id not in allowed_user_ids: + await query.answer("Access denied.", show_alert=True) + return ConversationHandler.END + return await func(update, context, *args, **kwargs) + return wrapper + + +def safe_delete_file(path: str) -> None: + """Delete a temporary file safely, logging errors without raising. + + :param path: Absolute or relative path to the file to delete. + """ + if not path: + return + try: + if os.path.exists(path): + os.remove(path) + logging.info("%s has been deleted.", path) + except Exception as e: + logging.error("Error deleting image file %s: %s", path, e) + + async def handle_api_error(update: Update, error_data: dict) -> bool: """Handle API error responses uniformly. Sends an appropriate error photo depending on error type. diff --git a/src/handlers/node.py b/src/handlers/node.py index 0c540c0..9966e9b 100644 --- a/src/handlers/node.py +++ b/src/handlers/node.py @@ -1,12 +1,14 @@ import os import logging -from datetime import datetime from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup from telegram.ext import CallbackContext, ConversationHandler from services.massa_rpc import get_addresses from services.docker_manager import start_docker_node, stop_docker_node, exec_massa_client -from handlers.common import auth_required, handle_api_error -from services.history import save_balance_history, get_entry_balance, get_entry_temperature, get_entry_ram +from handlers.common import auth_required, cb_auth_required, handle_api_error, safe_delete_file +from services.history import ( + save_balance_history, + make_time_key, build_balance_entry, format_history_entry, +) from services.plotting import create_png_plot, create_balance_history_plot, create_resources_plot from services.system_monitor import get_system_stats from config import ( @@ -18,6 +20,23 @@ ) +_DOCKER_MENU_TEXT = "🐳 Docker Node Management\nWhat do you want to do?" + + +def _build_docker_main_menu_markup() -> InlineKeyboardMarkup: + """Build the main Docker management inline keyboard.""" + keyboard = [ + [ + InlineKeyboardButton("▶️ Start", callback_data='docker_start'), + InlineKeyboardButton("⏹️ Stop", callback_data='docker_stop'), + ], + [ + InlineKeyboardButton("💻 Massa Client", callback_data='docker_massa'), + ], + ] + return InlineKeyboardMarkup(keyboard) + + def extract_address_data(json_data: dict): """ Extract useful JSON response data from get_address. @@ -69,18 +88,9 @@ async def node(update: Update, context: CallbackContext) -> None: await update.message.reply_text('Node status: ' + formatted_string) # Record current balance snapshot with timestamp, including system resources - now = datetime.now() - time_key = f"{now.year}/{now.month:02d}/{now.day:02d}-{now.hour:02d}:{now.minute:02d}" - system_stats = get_system_stats(logging) - temperature_avg = system_stats.get("temperature_avg") - ram_percent = system_stats.get("ram_percent") - - entry = {"balance": float(data[0])} - if temperature_avg is not None: - entry["temperature_avg"] = temperature_avg - if ram_percent is not None: - entry["ram_percent"] = ram_percent + time_key = make_time_key() + entry = build_balance_entry(float(data[0]), system_stats) lock = context.bot_data['balance_lock'] with lock: @@ -105,13 +115,7 @@ async def node(update: Update, context: CallbackContext) -> None: await update.message.reply_photo(photo=f'media/{PAT_FILE_NAME}') finally: # Always clean up the temporary chart image - if image_path: - try: - if os.path.exists(image_path): - os.remove(image_path) - logging.info(f"{image_path} has been deleted.") - except Exception as e: - logging.error(f"Error deleting image file {image_path}: {e}") + safe_delete_file(image_path) async def flush(update: Update, context: CallbackContext) -> int: @@ -152,16 +156,12 @@ async def flush(update: Update, context: CallbackContext) -> int: return FLUSH_CONFIRM_STATE +@cb_auth_required async def flush_confirm_yes(update: Update, context: CallbackContext) -> int: """Callback for flush 'Yes': clear both log file and balance history.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} confirmed flush with balance history clear.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: # Truncate the log file @@ -186,16 +186,12 @@ async def flush_confirm_yes(update: Update, context: CallbackContext) -> int: return ConversationHandler.END +@cb_auth_required async def flush_confirm_no(update: Update, context: CallbackContext) -> int: """Callback for flush 'No': clear only the log file, keep balance history.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} confirmed flush without balance history clear.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: # Truncate the log file only @@ -292,45 +288,25 @@ async def hist(update: Update, context: CallbackContext) -> int: finally: # Always clean up the temporary chart images for path in (image_path, resources_path): - if path: - try: - if os.path.exists(path): - os.remove(path) - logging.info(f"{path} has been deleted.") - except Exception as e: - logging.error(f"Error deleting image file {path}: {e}") + safe_delete_file(path) +@cb_auth_required async def hist_confirm_yes(update: Update, context: CallbackContext) -> int: """Callback for hist 'Yes': send the full balance history as text.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} confirmed hist with text summary.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) balance_history = context.bot_data['balance_history'] - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END - try: if not balance_history: await query.answer("Balance history is empty.", show_alert=True) return ConversationHandler.END # Format all history entries as a single text message - def _format_entry(time_key: str, value) -> str: - line = f"{time_key}: Balance {get_entry_balance(value):.2f}" - temp = get_entry_temperature(value) - ram = get_entry_ram(value) - if temp is not None: - line += f", Temp {temp:.1f}°C" - if ram is not None: - line += f", RAM {ram:.1f}%" - return line - tmp_string = "History\n" + "\n".join( - _format_entry(time_key, value) for time_key, value in balance_history.items() + format_history_entry(time_key, value) for time_key, value in balance_history.items() ) await query.edit_message_text(text="✓ Sending balance history...") @@ -344,16 +320,12 @@ def _format_entry(time_key: str, value) -> str: return ConversationHandler.END +@cb_auth_required async def hist_confirm_no(update: Update, context: CallbackContext) -> int: """Callback for hist 'No': dismiss the prompt.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} declined hist text summary.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: await query.edit_message_text(text="✓ Done.") @@ -378,37 +350,20 @@ async def docker(update: Update, context: CallbackContext) -> int: await update.message.reply_text("Access denied. You are not authorized.") return ConversationHandler.END - # Present inline keyboard: Start, Stop or Massa Client - keyboard = [ - [ - InlineKeyboardButton("▶️ Start", callback_data='docker_start'), - InlineKeyboardButton("⏹️ Stop", callback_data='docker_stop') - ], - [ - InlineKeyboardButton("💻 Massa Client", callback_data='docker_massa') - ] - ] - reply_markup = InlineKeyboardMarkup(keyboard) - await update.message.reply_text( - "🐳 Docker Node Management\n" - "What do you want to do?", - reply_markup=reply_markup + _DOCKER_MENU_TEXT, + reply_markup=_build_docker_main_menu_markup(), ) return DOCKER_MENU_STATE +@cb_auth_required async def docker_start(update: Update, context: CallbackContext) -> int: """Callback for docker 'Start': ask for confirmation.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} selected Start in docker menu.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END keyboard = [ [ @@ -427,16 +382,12 @@ async def docker_start(update: Update, context: CallbackContext) -> int: return DOCKER_START_CONFIRM_STATE +@cb_auth_required async def docker_stop(update: Update, context: CallbackContext) -> int: """Callback for docker 'Stop': ask for confirmation.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} selected Stop in docker menu.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END keyboard = [ [ @@ -455,16 +406,12 @@ async def docker_stop(update: Update, context: CallbackContext) -> int: return DOCKER_STOP_CONFIRM_STATE +@cb_auth_required async def docker_start_confirm(update: Update, context: CallbackContext) -> int: """Callback for docker start confirmation: execute docker start command.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} confirmed docker start.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: # Get container name from bot_data (must be set in main.py from topology.json) @@ -492,16 +439,12 @@ async def docker_start_confirm(update: Update, context: CallbackContext) -> int: return ConversationHandler.END +@cb_auth_required async def docker_stop_confirm(update: Update, context: CallbackContext) -> int: """Callback for docker stop confirmation: execute docker stop command.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} confirmed docker stop.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: # Get container name from bot_data (must be set in main.py from topology.json) @@ -529,16 +472,12 @@ async def docker_stop_confirm(update: Update, context: CallbackContext) -> int: return ConversationHandler.END +@cb_auth_required async def docker_cancel(update: Update, context: CallbackContext) -> int: """Callback for docker cancel: dismiss the action.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} cancelled docker action.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: await query.edit_message_text(text="❌ Action cancelled.") @@ -550,16 +489,12 @@ async def docker_cancel(update: Update, context: CallbackContext) -> int: return ConversationHandler.END +@cb_auth_required async def docker_massa(update: Update, context: CallbackContext) -> int: """Callback for docker 'Massa Client': show wallet_info / buy_rolls sub-menu.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} selected Massa Client in docker menu.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END keyboard = [ [ @@ -585,16 +520,12 @@ async def docker_massa(update: Update, context: CallbackContext) -> int: return DOCKER_MASSA_MENU_STATE +@cb_auth_required async def massa_wallet_info(update: Update, context: CallbackContext) -> int: """Callback for 'Wallet Info': execute wallet_info via massa-client.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} requested wallet_info.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: container_name = context.bot_data.get('docker_container_name') @@ -621,16 +552,12 @@ async def massa_wallet_info(update: Update, context: CallbackContext) -> int: return ConversationHandler.END +@cb_auth_required async def massa_buy_rolls_ask(update: Update, context: CallbackContext) -> int: """Callback for 'Buy Rolls': ask the user how many rolls to buy.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} selected Buy Rolls.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END await query.edit_message_text( text="🎲 Buy Rolls\n\n" @@ -642,16 +569,12 @@ async def massa_buy_rolls_ask(update: Update, context: CallbackContext) -> int: return DOCKER_BUYROLLS_INPUT_STATE +@cb_auth_required async def massa_sell_rolls_ask(update: Update, context: CallbackContext) -> int: """Callback for 'Sell Rolls': ask the user how many rolls to sell.""" query = update.callback_query user_id = str(query.from_user.id) logging.info(f'User {user_id} selected Sell Rolls.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END await query.edit_message_text( text="💸 Sell Rolls\n\n" @@ -663,8 +586,26 @@ async def massa_sell_rolls_ask(update: Update, context: CallbackContext) -> int: return DOCKER_SELLROLLS_INPUT_STATE -async def massa_buy_rolls_input(update: Update, context: CallbackContext) -> int: - """Handle text input for the number of rolls to buy.""" +async def _rolls_input_handler( + update: Update, + context: CallbackContext, + *, + action_label: str, + user_data_key: str, + confirm_callback: str, + input_state: int, + confirm_state: int, +) -> int: + """Shared input handler for buy/sell rolls conversation steps. + + Validates the user's roll count, stores it, and presents a confirmation keyboard. + + :param action_label: Human-readable action name (e.g. ``"buy_rolls"``). + :param user_data_key: Key to use in ``context.user_data`` for storing the count. + :param confirm_callback: ``callback_data`` for the 'Yes' confirmation button. + :param input_state: ConversationHandler state to return when input is invalid. + :param confirm_state: ConversationHandler state to return on valid input. + """ user_id = str(update.effective_user.id) allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) @@ -683,24 +624,24 @@ async def massa_buy_rolls_input(update: Update, context: CallbackContext) -> int await update.message.reply_text( "❌ Invalid number. Please send a positive integer (e.g. 1, 5, 10) or /docker to cancel." ) - return DOCKER_BUYROLLS_INPUT_STATE + return input_state # Store roll count for the confirmation step - context.user_data['buy_rolls_count'] = roll_count + context.user_data[user_data_key] = roll_count wallet_address = context.bot_data.get('massa_wallet_address', 'N/A') fee = context.bot_data.get('massa_buy_rolls_fee', 0.01) keyboard = [ [ - InlineKeyboardButton("Yes", callback_data='buyrolls_confirm'), - InlineKeyboardButton("No", callback_data='docker_cancel') + InlineKeyboardButton("Yes", callback_data=confirm_callback), + InlineKeyboardButton("No", callback_data='docker_cancel'), ] ] reply_markup = InlineKeyboardMarkup(keyboard) await update.message.reply_text( - f"⚠️ Confirm buy_rolls:\n\n" + f"⚠️ Confirm {action_label}:\n\n" f"Address: {wallet_address}\n" f"Rolls: {roll_count}\n" f"Fee: {fee}\n\n" @@ -708,176 +649,121 @@ async def massa_buy_rolls_input(update: Update, context: CallbackContext) -> int reply_markup=reply_markup ) - return DOCKER_BUYROLLS_CONFIRM_STATE - - -async def massa_buy_rolls_confirm(update: Update, context: CallbackContext) -> int: - """Callback for buy rolls confirmation: execute the buy_rolls command.""" - query = update.callback_query - user_id = str(query.from_user.id) - logging.info(f'User {user_id} confirmed buy_rolls.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) + return confirm_state - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END - try: - container_name = context.bot_data.get('docker_container_name') - password = context.bot_data.get('massa_client_password') - wallet_address = context.bot_data.get('massa_wallet_address') - fee = context.bot_data.get('massa_buy_rolls_fee', 0.01) - roll_count = context.user_data.get('buy_rolls_count', 0) - - if not all([container_name, password, wallet_address]) or roll_count <= 0: - await query.edit_message_text(text="❌ Error: Missing configuration or invalid roll count.") - await query.answer() - return ConversationHandler.END - - command = f"buy_rolls {wallet_address} {roll_count} {fee}" - - await query.edit_message_text(text=f"⏳ Executing buy_rolls ({roll_count} rolls)...") - await query.answer() - - result = exec_massa_client(logging, container_name, password, command) - - if result['status'] == 'ok': - output = result['output'] or 'Command executed (no output).' - await query.edit_message_text( - text=f"✅ buy_rolls executed:\n\n{output}" - ) - logging.info(f"User {user_id} bought {roll_count} rolls.") - else: - await query.edit_message_text(text=result['message']) - except Exception as e: - logging.error(f"Error executing buy_rolls: {e}") - await query.edit_message_text(text="❌ Error executing buy_rolls.") - - # Clean up user_data - context.user_data.pop('buy_rolls_count', None) - return ConversationHandler.END +async def massa_buy_rolls_input(update: Update, context: CallbackContext) -> int: + """Handle text input for the number of rolls to buy.""" + return await _rolls_input_handler( + update, context, + action_label="buy_rolls", + user_data_key='buy_rolls_count', + confirm_callback='buyrolls_confirm', + input_state=DOCKER_BUYROLLS_INPUT_STATE, + confirm_state=DOCKER_BUYROLLS_CONFIRM_STATE, + ) async def massa_sell_rolls_input(update: Update, context: CallbackContext) -> int: """Handle text input for the number of rolls to sell.""" - user_id = str(update.effective_user.id) - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await update.message.reply_text("Access denied.") - return ConversationHandler.END - - text = update.message.text.strip() - - # Validate input: must be a positive integer - try: - roll_count = int(text) - if roll_count <= 0: - raise ValueError - except ValueError: - await update.message.reply_text( - "❌ Invalid number. Please send a positive integer (e.g. 1, 5, 10) or /docker to cancel." - ) - return DOCKER_SELLROLLS_INPUT_STATE - - # Store roll count for the confirmation step - context.user_data['sell_rolls_count'] = roll_count - - wallet_address = context.bot_data.get('massa_wallet_address', 'N/A') - fee = context.bot_data.get('massa_buy_rolls_fee', 0.01) - - keyboard = [ - [ - InlineKeyboardButton("Yes", callback_data='sellrolls_confirm'), - InlineKeyboardButton("No", callback_data='docker_cancel') - ] - ] - reply_markup = InlineKeyboardMarkup(keyboard) - - await update.message.reply_text( - f"⚠️ Confirm sell_rolls:\n\n" - f"Address: {wallet_address}\n" - f"Rolls: {roll_count}\n" - f"Fee: {fee}\n\n" - f"Proceed?", - reply_markup=reply_markup + return await _rolls_input_handler( + update, context, + action_label="sell_rolls", + user_data_key='sell_rolls_count', + confirm_callback='sellrolls_confirm', + input_state=DOCKER_SELLROLLS_INPUT_STATE, + confirm_state=DOCKER_SELLROLLS_CONFIRM_STATE, ) - return DOCKER_SELLROLLS_CONFIRM_STATE +async def _rolls_exec_handler( + update: Update, + context: CallbackContext, + *, + command_name: str, + user_data_key: str, + action_verb: str, +) -> int: + """Shared execution handler for buy/sell rolls confirmation steps. -async def massa_sell_rolls_confirm(update: Update, context: CallbackContext) -> int: - """Callback for sell rolls confirmation: execute the sell_rolls command.""" + Retrieves stored roll count, executes the massa-client command, and cleans up. + + :param command_name: Massa-client command prefix (``"buy_rolls"`` or ``"sell_rolls"``). + :param user_data_key: Key in ``context.user_data`` holding the roll count. + :param action_verb: Past-tense verb for logging (``"bought"`` or ``"sold"``). + """ query = update.callback_query user_id = str(query.from_user.id) - logging.info(f'User {user_id} confirmed sell_rolls.') - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END try: container_name = context.bot_data.get('docker_container_name') password = context.bot_data.get('massa_client_password') wallet_address = context.bot_data.get('massa_wallet_address') fee = context.bot_data.get('massa_buy_rolls_fee', 0.01) - roll_count = context.user_data.get('sell_rolls_count', 0) + roll_count = context.user_data.get(user_data_key, 0) if not all([container_name, password, wallet_address]) or roll_count <= 0: await query.edit_message_text(text="❌ Error: Missing configuration or invalid roll count.") await query.answer() return ConversationHandler.END - command = f"sell_rolls {wallet_address} {roll_count} {fee}" + command = f"{command_name} {wallet_address} {roll_count} {fee}" - await query.edit_message_text(text=f"⏳ Executing sell_rolls ({roll_count} rolls)...") + await query.edit_message_text(text=f"⏳ Executing {command_name} ({roll_count} rolls)...") await query.answer() result = exec_massa_client(logging, container_name, password, command) if result['status'] == 'ok': output = result['output'] or 'Command executed (no output).' - await query.edit_message_text( - text=f"✅ sell_rolls executed:\n\n{output}" - ) - logging.info(f"User {user_id} sold {roll_count} rolls.") + await query.edit_message_text(text=f"✅ {command_name} executed:\n\n{output}") + logging.info(f"User {user_id} {action_verb} {roll_count} rolls.") else: await query.edit_message_text(text=result['message']) except Exception as e: - logging.error(f"Error executing sell_rolls: {e}") - await query.edit_message_text(text="❌ Error executing sell_rolls.") + logging.error(f"Error executing {command_name}: {e}") + await query.edit_message_text(text=f"❌ Error executing {command_name}.") # Clean up user_data - context.user_data.pop('sell_rolls_count', None) + context.user_data.pop(user_data_key, None) return ConversationHandler.END +@cb_auth_required +async def massa_buy_rolls_confirm(update: Update, context: CallbackContext) -> int: + """Callback for buy rolls confirmation: execute the buy_rolls command.""" + user_id = str(update.callback_query.from_user.id) + logging.info(f'User {user_id} confirmed buy_rolls.') + return await _rolls_exec_handler( + update, context, + command_name='buy_rolls', + user_data_key='buy_rolls_count', + action_verb='bought', + ) + + +@cb_auth_required +async def massa_sell_rolls_confirm(update: Update, context: CallbackContext) -> int: + """Callback for sell rolls confirmation: execute the sell_rolls command.""" + user_id = str(update.callback_query.from_user.id) + logging.info(f'User {user_id} confirmed sell_rolls.') + return await _rolls_exec_handler( + update, context, + command_name='sell_rolls', + user_data_key='sell_rolls_count', + action_verb='sold', + ) + + +@cb_auth_required async def massa_back(update: Update, context: CallbackContext) -> int: """Callback to go back to the main docker menu.""" query = update.callback_query user_id = str(query.from_user.id) - allowed_user_ids = context.bot_data.get('allowed_user_ids', set()) - - if user_id not in allowed_user_ids: - await query.answer("Access denied.", show_alert=True) - return ConversationHandler.END - - keyboard = [ - [ - InlineKeyboardButton("▶️ Start", callback_data='docker_start'), - InlineKeyboardButton("⏹️ Stop", callback_data='docker_stop') - ], - [ - InlineKeyboardButton("💻 Massa Client", callback_data='docker_massa') - ] - ] - reply_markup = InlineKeyboardMarkup(keyboard) await query.edit_message_text( - text="🐳 Docker Node Management\n" - "What do you want to do?", - reply_markup=reply_markup + text=_DOCKER_MENU_TEXT, + reply_markup=_build_docker_main_menu_markup(), ) await query.answer() diff --git a/src/handlers/scheduler.py b/src/handlers/scheduler.py index ba91825..d1ad32d 100644 --- a/src/handlers/scheduler.py +++ b/src/handlers/scheduler.py @@ -7,7 +7,11 @@ from services.massa_rpc import get_addresses from services.system_monitor import get_system_stats from handlers.node import extract_address_data -from services.history import save_balance_history, filter_last_24h, filter_since_midnight, get_entry_balance, get_entry_temperature, get_entry_ram +from services.history import ( + save_balance_history, filter_last_24h, filter_since_midnight, + get_entry_balance, get_entry_temperature, + make_time_key, build_balance_entry, format_history_entry, +) from config import ( JOB_SCHED_NAME, NODE_IS_DOWN, NODE_IS_UP, TIMEOUT_NAME, TIMEOUT_FIRE_NAME, @@ -143,19 +147,12 @@ async def periodic_node_ping(application: Application) -> None: # Record current balance snapshot with timestamp, including system resources now = datetime.now() - hour, minute, day, month, year = now.hour, now.minute, now.day, now.month, now.year - current_time_key = f"{year}/{month:02d}/{day:02d}-{hour:02d}:{minute:02d}" + hour = now.hour + current_time_key = make_time_key(now) # Collect CPU temperature and RAM usage system_stats = get_system_stats(logging) - temperature_avg = system_stats.get("temperature_avg") - ram_percent = system_stats.get("ram_percent") - - entry = {"balance": float(data[0])} - if temperature_avg is not None: - entry["temperature_avg"] = temperature_avg - if ram_percent is not None: - entry["ram_percent"] = ram_percent + entry = build_balance_entry(float(data[0]), system_stats) lock = application.bot_data.get('balance_lock') if lock: @@ -229,9 +226,7 @@ async def periodic_node_ping(application: Application) -> None: f"📊 Last 24h History:\n" f"{'─' * 40}\n" + ("\n".join( - f"{timestamp}: Balance {get_entry_balance(v):.2f}" - + (f", Temp {get_entry_temperature(v):.1f}°C" if get_entry_temperature(v) is not None else "") - + (f", RAM {get_entry_ram(v):.1f}%" if get_entry_ram(v) is not None else "") + format_history_entry(timestamp, v) for timestamp, v in recent_history.items() ) if recent_history else "No data in the last 24h.") ) diff --git a/src/services/history.py b/src/services/history.py index 8674189..15c3106 100644 --- a/src/services/history.py +++ b/src/services/history.py @@ -8,6 +8,54 @@ BALANCE_HISTORY_FILE = 'config/balance_history.json' +def make_time_key(dt: datetime = None) -> str: + """Return a ``YYYY/MM/DD-HH:MM`` time key suitable for balance history. + + :param dt: Datetime to format; defaults to now. + :return: Formatted time key string. + """ + if dt is None: + dt = datetime.now() + return f"{dt.year}/{dt.month:02d}/{dt.day:02d}-{dt.hour:02d}:{dt.minute:02d}" + + +def build_balance_entry(balance: float, system_stats: dict) -> dict: + """Build a balance history entry dict from a balance and system stats. + + Includes ``temperature_avg`` and ``ram_percent`` when they are present + in *system_stats*. + + :param balance: Current node balance. + :param system_stats: Dict returned by ``get_system_stats``. + :return: Entry dict ready to store in balance history. + """ + entry: dict = {"balance": balance} + temperature_avg = system_stats.get("temperature_avg") + ram_percent = system_stats.get("ram_percent") + if temperature_avg is not None: + entry["temperature_avg"] = temperature_avg + if ram_percent is not None: + entry["ram_percent"] = ram_percent + return entry + + +def format_history_entry(time_key: str, value) -> str: + """Format a single history entry as a human-readable string. + + :param time_key: Timestamp key (e.g. ``"2025/03/14-07:00"``). + :param value: History entry value (str or dict). + :return: Formatted string like ``"2025/03/14-07:00: Balance 1234.56, Temp 42.0°C, RAM 63.5%"``. + """ + line = f"{time_key}: Balance {get_entry_balance(value):.2f}" + temp = get_entry_temperature(value) + ram = get_entry_ram(value) + if temp is not None: + line += f", Temp {temp:.1f}°C" + if ram is not None: + line += f", RAM {ram:.1f}%" + return line + + def get_entry_balance(value) -> float: """Extract the balance from a history entry. diff --git a/tests/test_common_helpers.py b/tests/test_common_helpers.py new file mode 100644 index 0000000..b5520f5 --- /dev/null +++ b/tests/test_common_helpers.py @@ -0,0 +1,258 @@ +"""Tests for the new factorized helpers in handlers/common.py. + +Covers: + - cb_auth_required (authorized + unauthorized paths) + - safe_delete_file (existing file, missing file, None/empty path) + - auth_required (authorized + unauthorized paths) + - handle_api_error (no error, generic error, timeout error) +""" +import sys +import os +import tempfile +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +# Ensure src/ is on the path so imports work without installation +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) + +from handlers.common import cb_auth_required, safe_delete_file, auth_required, handle_api_error +from telegram.ext import ConversationHandler + + +# --------------------------------------------------------------------------- +# Helpers to build minimal Telegram mock objects +# --------------------------------------------------------------------------- + +def _make_context(allowed_user_ids=None): + """Build a minimal CallbackContext mock.""" + ctx = MagicMock() + ctx.bot_data = {"allowed_user_ids": allowed_user_ids or set()} + return ctx + + +def _make_message_update(user_id: str): + """Build a minimal Update mock for message-based handlers.""" + update = MagicMock() + update.effective_user.id = user_id + update.message.reply_text = AsyncMock() + return update + + +def _make_callback_update(user_id: str): + """Build a minimal Update mock for callback query handlers.""" + update = MagicMock() + update.callback_query.from_user.id = user_id + update.callback_query.answer = AsyncMock() + return update + + +# --------------------------------------------------------------------------- +# cb_auth_required +# --------------------------------------------------------------------------- + +class TestCbAuthRequired: + @pytest.mark.asyncio + async def test_allows_authorized_user(self): + allowed = {"42"} + ctx = _make_context(allowed) + update = _make_callback_update("42") + + @cb_auth_required + async def handler(update, context): + return "ok" + + result = await handler(update, ctx) + assert result == "ok" + + @pytest.mark.asyncio + async def test_blocks_unauthorized_user(self): + ctx = _make_context({"99"}) # user 42 is NOT in the list + update = _make_callback_update("42") + + @cb_auth_required + async def handler(update, context): + return "ok" + + result = await handler(update, ctx) + assert result == ConversationHandler.END + + @pytest.mark.asyncio + async def test_unauthorized_calls_query_answer(self): + ctx = _make_context(set()) + update = _make_callback_update("42") + + @cb_auth_required + async def handler(update, context): + return "ok" + + await handler(update, ctx) + update.callback_query.answer.assert_awaited_once() + + @pytest.mark.asyncio + async def test_empty_allowlist_blocks_all(self): + ctx = _make_context(set()) + update = _make_callback_update("1") + + @cb_auth_required + async def handler(update, context): + return "reached" + + result = await handler(update, ctx) + assert result == ConversationHandler.END + + @pytest.mark.asyncio + async def test_preserves_function_name(self): + @cb_auth_required + async def my_handler(update, context): + pass + + assert my_handler.__name__ == "my_handler" + + @pytest.mark.asyncio + async def test_authorized_user_receives_context_args(self): + """Authorized handler should receive update and context unmodified.""" + ctx = _make_context({"5"}) + update = _make_callback_update("5") + received = {} + + @cb_auth_required + async def handler(update, context): + received["update"] = update + received["context"] = context + + await handler(update, ctx) + assert received["update"] is update + assert received["context"] is ctx + + +# --------------------------------------------------------------------------- +# auth_required +# --------------------------------------------------------------------------- + +class TestAuthRequired: + @pytest.mark.asyncio + async def test_allows_authorized_user(self): + ctx = _make_context({"7"}) + update = _make_message_update("7") + + @auth_required + async def handler(update, context): + return "ok" + + result = await handler(update, ctx) + assert result == "ok" + + @pytest.mark.asyncio + async def test_blocks_unauthorized_user(self): + ctx = _make_context({"99"}) + update = _make_message_update("1") + + @auth_required + async def handler(update, context): + return "reached" + + result = await handler(update, ctx) + assert result is None # auth_required returns None on denial + + @pytest.mark.asyncio + async def test_unauthorized_replies_with_message(self): + ctx = _make_context(set()) + update = _make_message_update("1") + + @auth_required + async def handler(update, context): + pass + + await handler(update, ctx) + update.message.reply_text.assert_awaited_once() + + @pytest.mark.asyncio + async def test_preserves_function_name(self): + @auth_required + async def my_cmd(update, context): + pass + + assert my_cmd.__name__ == "my_cmd" + + +# --------------------------------------------------------------------------- +# safe_delete_file +# --------------------------------------------------------------------------- + +class TestSafeDeleteFile: + def test_deletes_existing_file(self, tmp_path): + f = tmp_path / "chart.png" + f.write_bytes(b"data") + safe_delete_file(str(f)) + assert not f.exists() + + def test_does_not_raise_for_missing_file(self): + safe_delete_file("/tmp/nonexistent_robbi_test_xyz.png") # must not raise + + def test_does_not_raise_for_none(self): + safe_delete_file(None) # must not raise + + def test_does_not_raise_for_empty_string(self): + safe_delete_file("") # must not raise + + def test_logs_deletion(self, tmp_path, caplog): + import logging + f = tmp_path / "tmp.png" + f.write_bytes(b"x") + with caplog.at_level(logging.INFO): + safe_delete_file(str(f)) + assert any("deleted" in record.message.lower() for record in caplog.records) + + def test_no_error_logged_for_missing_file(self, caplog): + import logging + with caplog.at_level(logging.ERROR): + safe_delete_file("/tmp/totally_absent_robbi_xyz.png") + assert not caplog.records + + +# --------------------------------------------------------------------------- +# handle_api_error +# --------------------------------------------------------------------------- + +class TestHandleApiError: + @pytest.mark.asyncio + async def test_returns_false_when_no_error(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + result = await handle_api_error(update, {"result": [{"balance": "100"}]}) + assert result is False + + @pytest.mark.asyncio + async def test_returns_true_on_generic_error(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + result = await handle_api_error(update, {"error": "Some API error"}) + assert result is True + + @pytest.mark.asyncio + async def test_returns_true_on_timeout_error(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + result = await handle_api_error(update, {"error": "Request timed out"}) + assert result is True + + @pytest.mark.asyncio + async def test_sends_photo_on_generic_error(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + await handle_api_error(update, {"error": "Connection refused"}) + update.message.reply_photo.assert_awaited_once() + + @pytest.mark.asyncio + async def test_sends_photo_on_timeout_error(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + await handle_api_error(update, {"error": "timed out"}) + update.message.reply_photo.assert_awaited_once() + + @pytest.mark.asyncio + async def test_empty_dict_returns_false(self): + update = _make_message_update("1") + update.message.reply_photo = AsyncMock() + result = await handle_api_error(update, {}) + assert result is False diff --git a/tests/test_history_helpers.py b/tests/test_history_helpers.py new file mode 100644 index 0000000..2a3efa2 --- /dev/null +++ b/tests/test_history_helpers.py @@ -0,0 +1,148 @@ +"""Tests for the new factorized helpers in services/history.py. + +Covers: + - make_time_key + - build_balance_entry + - format_history_entry +""" +import sys +import os +import pytest +from datetime import datetime + +# Ensure src/ is on the path so imports work without installation +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) + +from services.history import make_time_key, build_balance_entry, format_history_entry + + +# --------------------------------------------------------------------------- +# make_time_key +# --------------------------------------------------------------------------- + +class TestMakeTimeKey: + def test_formats_datetime_correctly(self): + dt = datetime(2025, 3, 14, 7, 5) + assert make_time_key(dt) == "2025/03/14-07:05" + + def test_zero_pads_month_day_hour_minute(self): + dt = datetime(2025, 1, 2, 3, 4) + assert make_time_key(dt) == "2025/01/02-03:04" + + def test_end_of_year(self): + dt = datetime(2024, 12, 31, 23, 59) + assert make_time_key(dt) == "2024/12/31-23:59" + + def test_default_uses_current_time(self): + before = datetime.now() + key = make_time_key() + after = datetime.now() + # Key must be parseable back to a datetime + dt = datetime.strptime(key, "%Y/%m/%d-%H:%M") + # The key's datetime must lie within [before, after] (minute precision) + assert dt.year == before.year + assert dt.month == before.month + assert dt.day == before.day + assert dt.hour == before.hour + + def test_returns_string(self): + assert isinstance(make_time_key(datetime(2025, 6, 15, 12, 0)), str) + + +# --------------------------------------------------------------------------- +# build_balance_entry +# --------------------------------------------------------------------------- + +class TestBuildBalanceEntry: + def test_balance_only_when_stats_empty(self): + entry = build_balance_entry(1234.56, {}) + assert entry == {"balance": 1234.56} + + def test_includes_temperature_when_present(self): + entry = build_balance_entry(100.0, {"temperature_avg": 42.5}) + assert entry == {"balance": 100.0, "temperature_avg": 42.5} + + def test_includes_ram_when_present(self): + entry = build_balance_entry(100.0, {"ram_percent": 63.2}) + assert entry == {"balance": 100.0, "ram_percent": 63.2} + + def test_includes_both_temperature_and_ram(self): + entry = build_balance_entry(500.0, {"temperature_avg": 55.0, "ram_percent": 80.1}) + assert entry == {"balance": 500.0, "temperature_avg": 55.0, "ram_percent": 80.1} + + def test_excludes_temperature_when_none(self): + entry = build_balance_entry(100.0, {"temperature_avg": None, "ram_percent": 50.0}) + assert "temperature_avg" not in entry + assert entry["ram_percent"] == 50.0 + + def test_excludes_ram_when_none(self): + entry = build_balance_entry(100.0, {"temperature_avg": 40.0, "ram_percent": None}) + assert "ram_percent" not in entry + assert entry["temperature_avg"] == 40.0 + + def test_ignores_unrelated_stats_keys(self): + entry = build_balance_entry(200.0, {"cpu_percent": 10.0, "other": "ignored"}) + assert entry == {"balance": 200.0} + + def test_balance_stored_as_float(self): + entry = build_balance_entry(0.0, {}) + assert isinstance(entry["balance"], float) + + def test_zero_balance(self): + entry = build_balance_entry(0.0, {}) + assert entry["balance"] == 0.0 + + +# --------------------------------------------------------------------------- +# format_history_entry +# --------------------------------------------------------------------------- + +class TestFormatHistoryEntry: + def test_dict_entry_with_all_fields(self): + value = {"balance": 1234.56, "temperature_avg": 42.1, "ram_percent": 63.5} + result = format_history_entry("2025/03/14-07:05", value) + assert result == "2025/03/14-07:05: Balance 1234.56, Temp 42.1°C, RAM 63.5%" + + def test_dict_entry_balance_only(self): + value = {"balance": 500.0} + result = format_history_entry("2025/03/14-08:00", value) + assert result == "2025/03/14-08:00: Balance 500.00" + + def test_dict_entry_with_temperature_no_ram(self): + value = {"balance": 200.0, "temperature_avg": 55.0} + result = format_history_entry("2025/03/14-09:00", value) + assert result == "2025/03/14-09:00: Balance 200.00, Temp 55.0°C" + + def test_dict_entry_with_ram_no_temperature(self): + value = {"balance": 300.0, "ram_percent": 75.3} + result = format_history_entry("2025/03/14-10:00", value) + assert result == "2025/03/14-10:00: Balance 300.00, RAM 75.3%" + + def test_legacy_string_entry(self): + # Legacy format: "Balance: 1234.56" + result = format_history_entry("14/03-07:05", "Balance: 1234.56") + assert result == "14/03-07:05: Balance 1234.56" + + def test_zero_balance(self): + value = {"balance": 0.0} + result = format_history_entry("2025/01/01-00:00", value) + assert result == "2025/01/01-00:00: Balance 0.00" + + def test_balance_rounds_to_two_decimals(self): + value = {"balance": 1234.5678} + result = format_history_entry("2025/03/14-07:05", value) + assert "Balance 1234.57" in result + + def test_temperature_rounds_to_one_decimal(self): + value = {"balance": 100.0, "temperature_avg": 42.16} + result = format_history_entry("2025/03/14-07:05", value) + assert "Temp 42.2°C" in result + + def test_ram_rounds_to_one_decimal(self): + value = {"balance": 100.0, "ram_percent": 63.16} + result = format_history_entry("2025/03/14-07:05", value) + assert "RAM 63.2%" in result + + def test_returns_string(self): + result = format_history_entry("2025/03/14-07:05", {"balance": 1.0}) + assert isinstance(result, str)