pumpfun-bonkfun-bot/src/cli.py

#!/usr/bin/env python3
"""
Command-line interface for the pump.fun trading bot.
"""

import argparse
import asyncio
import os
import sys

from dotenv import load_dotenv

import config
from trading.trader import PumpTrader
from utils.logger import get_logger, setup_file_logging

load_dotenv()

logger = get_logger(__name__)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments.

    Returns:
        Parsed arguments
    """
    parser = argparse.ArgumentParser(description="Trade tokens on pump.fun.")
    parser.add_argument(
        "--yolo", action="store_true", help="Run in YOLO mode (continuous trading)"
    )
    parser.add_argument(
        "--match",
        type=str,
        help="Only trade tokens with names or symbols matching this string",
    )
    parser.add_argument(
        "--bro", type=str, help="Only trade tokens created by this user address"
    )
    parser.add_argument(
        "--marry", action="store_true", help="Only buy tokens, skip selling"
    )
    parser.add_argument(
        "--amount",
        type=float,
        help=f"Amount of SOL to spend on each buy (default: {config.BUY_AMOUNT})",
    )
    parser.add_argument(
        "--buy-slippage",
        type=float,
        help=f"Buy slippage tolerance (default: {config.BUY_SLIPPAGE})",
    )
    parser.add_argument(
        "--sell-slippage",
        type=float,
        help=f"Sell slippage tolerance (default: {config.SELL_SLIPPAGE})",
    )

    return parser.parse_args()


async def main() -> None:
    """Main entry point for the CLI."""
    setup_file_logging("pump_trading.log")

    args = parse_args()

    # Get configuration values, preferring command line args over config.py
    rpc_endpoint: str | None = os.environ.get("SOLANA_NODE_RPC_ENDPOINT")
    wss_endpoint: str | None = os.environ.get("SOLANA_NODE_WSS_ENDPOINT")
    private_key: str | None = os.environ.get("SOLANA_PRIVATE_KEY")
    geyser_endpoint: str | None = os.environ.get("GEYSER_ENDPOINT")
    geyser_api_token: str | None = os.environ.get("GEYSER_API_TOKEN")


    # Validate configuration values
    if not rpc_endpoint or not rpc_endpoint.startswith(("http://", "https://")):
        logger.error("Invalid RPC endpoint. Must start with http:// or https://")
        sys.exit(1)

    if not wss_endpoint or not wss_endpoint.startswith(("ws://", "wss://")):
        logger.error("Invalid WebSocket endpoint. Must start with ws:// or wss://")
        sys.exit(1)

    if not private_key or len(private_key) < 80:
        logger.error("Invalid private key. Key appears to be missing or too short")
        sys.exit(1)

    if config.LISTENER_TYPE.lower() == "geyser":
        if not geyser_endpoint:
            logger.error("GEYSER_ENDPOINT environment variable is not set")
            sys.exit(1)
        if not geyser_api_token:
            logger.error("GEYSER_API_TOKEN environment variable is not set")
            sys.exit(1)

    # Get trading parameters
    buy_amount: float = args.amount if args.amount is not None else config.BUY_AMOUNT
    buy_slippage: float = (
        args.buy_slippage if args.buy_slippage is not None else config.BUY_SLIPPAGE
    )
    sell_slippage: float = (
        args.sell_slippage if args.sell_slippage is not None else config.SELL_SLIPPAGE
    )

    trader: PumpTrader = PumpTrader(
        rpc_endpoint=rpc_endpoint,  # type: ignore
        wss_endpoint=wss_endpoint,  # type: ignore
        private_key=private_key,
        buy_amount=buy_amount,
        buy_slippage=buy_slippage,
        sell_slippage=sell_slippage,
        max_retries=config.MAX_RETRIES,
        listener_type=config.LISTENER_TYPE,
        geyser_endpoint=geyser_endpoint,
        geyser_api_token=geyser_api_token,  
    )

    try:
        await trader.start(
            match_string=args.match,
            bro_address=args.bro,
            marry_mode=args.marry,
            yolo_mode=args.yolo,
        )
    except KeyboardInterrupt:
        logger.info("Trading stopped by user")
    except Exception as e:
        logger.error(f"Trading stopped due to error: {e!s}")
    finally:
        try:
            await trader.solana_client.close()
        except Exception:
            pass


def sync_main() -> None:
    asyncio.run(main())


if __name__ == "__main__":
    asyncio.run(main())