#!/bin/bash

# 交互式实时命令监控脚本 - 清理安装版
# 版本: 4.3

set -e

### 配置区域 ###
LOG_DIR="/root/command_monitor_logs"
SCRIPT_NAME="command_monitor.sh"
SCRIPT_PATH="/usr/local/bin/$SCRIPT_NAME"

### 颜色定义 ###
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

### 全局变量 ###
CURRENT_LOG=""
LATEST_LOG=""
DAEMON_MODE=false

# 清理所有残留文件和进程
cleanup_all() {
    echo -e "${YELLOW}正在清理所有残留...${NC}"
    
    # 停止所有监控进程
    echo -e "${BLUE}停止监控进程...${NC}"
    pkill -f "monitor_history_files" 2>/dev/null || true
    pkill -f "command_monitor" 2>/dev/null || true
    
    # 删除PID文件
    rm -f "/tmp/command_monitor.pid" 2>/dev/null || true
    rm -f "/tmp/last_size_"* 2>/dev/null || true
    
    # 删除缓存文件
    rm -f "/tmp/ip_geo_cache.txt" 2>/dev/null || true
    
    # 删除旧的自启动配置
    echo -e "${BLUE}清理自启动配置...${NC}"
    rm -f /etc/profile.d/command_monitor_auto.sh 2>/dev/null || true
    rm -f /etc/systemd/system/command-monitor.service 2>/dev/null || true
    systemctl daemon-reload 2>/dev/null || true
    
    # 清理rc.local中的配置
    if [ -f "/etc/rc.local" ]; then
        sed -i '/command_monitor.sh/d' "/etc/rc.local" 2>/dev/null || true
    fi
    if [ -f "/etc/rc.d/rc.local" ]; then
        sed -i '/command_monitor.sh/d' "/etc/rc.d/rc.local" 2>/dev/null || true
    fi
    
    # 清理日志文件（保留最新的3个）
    echo -e "${BLUE}清理日志文件...${NC}"
    if [ -d "$LOG_DIR" ]; then
        # 保留最新的3个日志文件
        ls -t "$LOG_DIR"/*.log 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true
        echo -e "${GREEN}日志文件清理完成${NC}"
    fi
    
    # 删除旧版本脚本
    rm -f "/usr/local/bin/command_monitor" 2>/dev/null || true
    rm -f "/usr/bin/command_monitor" 2>/dev/null || true
    
    echo -e "${GREEN}✓ 所有残留已清理完成${NC}"
}

# 检查磁盘空间
check_disk_space() {
    local available_space=$(df / | awk 'NR==2{print $4}')
    if [ "$available_space" -lt 10240 ]; then  # 小于10MB
        echo -e "${RED}磁盘空间不足，可用空间: ${available_space}KB${NC}"
        echo -e "${YELLOW}建议清理磁盘空间后再继续${NC}"
        return 1
    fi
    echo -e "${GREEN}磁盘空间充足: ${available_space}KB${NC}"
    return 0
}

# 获取时间戳
timestamp() {
    date '+%Y-%m-%d %H:%M:%S'
}

# 日志函数
log_message() {
    local level="$1"
    local message="$2"
    local color="$GREEN"

    case "$level" in
        "ERROR") color="$RED" ;;
        "WARN") color="$YELLOW" ;;
        "INFO") color="$BLUE" ;;
        "SUCCESS") color="$GREEN" ;;
        "COMMAND") color="$CYAN" ;;
    esac

    if [ "$DAEMON_MODE" = true ]; then
        echo -e "${color}[$(timestamp)] [$level] $message${NC}" >> "$CURRENT_LOG"
    else
        echo -e "${color}[$(timestamp)] [$level] $message${NC}" | tee -a "$CURRENT_LOG"
    fi
}

# 获取IP地理位置
get_ip_location() {
    local ip="$1"
    
    if [[ "$ip" == "192.168."* ]] || [[ "$ip" == "10."* ]] || [[ "$ip" == "172."* ]]; then
        echo "内网IP"
        return 0
    fi
    
    if [[ "$ip" == "127.0.0.1" ]] || [[ "$ip" == "localhost" ]] || [[ "$ip" == "unknown" ]]; then
        echo "本机"
        return 0
    fi
    
    local location_info=$(curl -s -m 3 "http://ip-api.com/json/$ip?fields=status,country,regionName,city,isp" 2>/dev/null || true)
    if [ -n "$location_info" ] && echo "$location_info" | grep -q '"status":"success"'; then
        local country=$(echo "$location_info" | grep -o '"country":"[^"]*"' | cut -d'"' -f4)
        local region=$(echo "$location_info" | grep -o '"regionName":"[^"]*"' | cut -d'"' -f4)
        local city=$(echo "$location_info" | grep -o '"city":"[^"]*"' | cut -d'"' -f4)
        local isp=$(echo "$location_info" | grep -o '"isp":"[^"]*"' | cut -d'"' -f4)
        
        if [ -n "$country" ]; then
            local result=""
            [ -n "$country" ] && result="$country"
            [ -n "$region" ] && result="$result-$region"
            [ -n "$city" ] && result="$result-$city"
            [ -n "$isp" ] && result="$result($isp)"
            echo "$result"
            return 0
        fi
    fi
    
    echo "未知位置"
}

# 获取客户端IP
get_client_ip() {
    local ip="unknown"
    
    if [ -n "$SSH_CLIENT" ]; then
        ip=$(echo "$SSH_CLIENT" | awk '{print $1}')
    elif [ -n "$SSH_CONNECTION" ]; then
        ip=$(echo "$SSH_CONNECTION" | awk '{print $1}')
    else
        ip=$(who -m 2>/dev/null | awk '{print $5}' | sed 's/[()]//g' | head -1)
        if [[ "$ip" == ":0" ]] || [[ "$ip" == ":1" ]] || [[ -z "$ip" ]]; then
            ip="localhost"
        fi
    fi
    
    echo "$ip"
}

# 初始化日志系统
init_log_system() {
    mkdir -p "$LOG_DIR"
    
    local client_ip=$(get_client_ip | sed 's/\./_/g')
    local log_date=$(date '+%Y%m%d_%H%M%S')
    CURRENT_LOG="$LOG_DIR/monitor_${client_ip}_${log_date}.log"
    LATEST_LOG="$LOG_DIR/latest.log"
    
    ln -sf "$CURRENT_LOG" "$LATEST_LOG" 2>/dev/null || true
}

# 配置实时history记录
configure_realtime_history() {
    for user_dir in /home/* /root; do
        if [ -d "$user_dir" ]; then
            local user=$(basename "$user_dir")
            local bashrc="$user_dir/.bashrc"
            
            if [ -f "$bashrc" ] && ! grep -q "REAL_TIME_HISTORY" "$bashrc"; then
                cat >> "$bashrc" << 'EOF'

# REAL_TIME_HISTORY - 实时命令记录配置
export PROMPT_COMMAND='history -a; history -c; history -r'
export HISTTIMEFORMAT='%F %T '
export HISTSIZE=10000
export HISTFILESIZE=20000
shopt -s histappend
export HISTCONTROL=ignoredups:erasedups

EOF
            fi
        fi
    done
}

# 监控命令历史文件变化
monitor_history_files() {
    declare -A file_sizes
    declare -A last_commands
    
    # 初始化文件大小
    for user_dir in /home/* /root; do
        if [ -d "$user_dir" ]; then
            local user=$(basename "$user_dir")
            local history_file="$user_dir/.bash_history"
            if [ -f "$history_file" ]; then
                file_sizes["$user"]=$(stat -c%s "$history_file" 2>/dev/null || echo 0)
                last_commands["$user"]=""
            fi
        fi
    done
    
    # 监控循环
    while true; do
        for user_dir in /home/* /root; do
            if [ -d "$user_dir" ]; then
                local user=$(basename "$user_dir")
                local history_file="$user_dir/.bash_history"
                
                if [ -f "$history_file" ]; then
                    local current_size=$(stat -c%s "$history_file" 2>/dev/null || echo 0)
                    local last_size=${file_sizes["$user"]}
                    
                    if [ "$current_size" -gt "$last_size" ]; then
                        local new_content=$(tail -n 1 "$history_file" 2>/dev/null | tr -d '\0')
                        
                        if [ -n "$new_content" ]; then
                            local line=$(echo "$new_content" | sed 's/^[ \t]*//;s/[ \t]*$//')
                            
                            # 过滤无用命令
                            if [ -n "$line" ] && [ "${#line}" -gt 1 ] && \
                               [[ ! "$line" =~ ^(ls|cd|pwd|ll|la|history|exit|clear|echo|date|whoami)$ ]] && \
                               [[ ! "$line" =~ ^#[0-9]+$ ]] && \
                               [ "$line" != "${last_commands["$user"]}" ]; then
                                
                                local client_ip=$(get_client_ip)
                                local location_info=$(get_ip_location "$client_ip")
                                
                                log_message "COMMAND" "用户: $user | 命令: $line | 来源: $client_ip [$location_info]"
                                last_commands["$user"]="$line"
                            fi
                        fi
                        
                        file_sizes["$user"]=$current_size
                    fi
                fi
            fi
        done
        sleep 1
    done
}

# 安装脚本到系统（优化版）
install_script() {
    echo -e "${GREEN}正在安装脚本到系统...${NC}"
    
    # 检查磁盘空间
    if ! check_disk_space; then
        echo -e "${RED}磁盘空间不足，安装中止${NC}"
        return 1
    fi
    
    # 先清理残留
    cleanup_all
    
    # 创建简化版脚本
    cat > "$SCRIPT_PATH" << 'EOF'
#!/bin/bash
# 简化版命令监控脚本

LOG_DIR="/root/command_monitor_logs"
PID_FILE="/tmp/command_monitor.pid"
LOG_FILE="$LOG_DIR/command_monitor.log"

start_daemon() {
    mkdir -p "$LOG_DIR"
    
    # 后台监控进程
    (
        while true; do
            for user_dir in /home/* /root; do
                [ -d "$user_dir" ] || continue
                user=$(basename "$user_dir")
                history_file="$user_dir/.bash_history"
                [ -f "$history_file" ] || continue
                
                current_size=$(stat -c%s "$history_file" 2>/dev/null || echo 0)
                last_size_file="/tmp/last_size_$user"
                last_size=$(cat "$last_size_file" 2>/dev/null || echo 0)
                
                if [ "$current_size" -gt "$last_size" ]; then
                    new_command=$(tail -n 1 "$history_file" 2>/dev/null | sed 's/^[ \t]*//;s/[ \t]*$//')
                    if [ -n "$new_command" ] && [ ${#new_command} -gt 2 ] && [[ ! "$new_command" =~ ^(ls|cd|pwd|history|exit|clear)$ ]]; then
                        client_ip=$(who -m 2>/dev/null | awk '{print $5}' | sed 's/[()]//g' | head -1)
                        echo "[$(date "+%Y-%m-%d %H:%M:%S")] 用户: $user | 命令: $new_command | 来源: $client_ip" >> "$LOG_FILE"
                    fi
                    echo "$current_size" > "$last_size_file"
                fi
            done
            sleep 1
        done
    ) &
    
    echo $! > "$PID_FILE"
    echo "监控已启动 (PID: $!)"
}

stop_daemon() {
    if [ -f "$PID_FILE" ]; then
        pid=$(cat "$PID_FILE")
        kill "$pid" 2>/dev/null && echo "监控已停止" || echo "停止失败"
        rm -f "$PID_FILE"
        rm -f /tmp/last_size_* 2>/dev/null
    else
        echo "监控未运行"
    fi
}

case "$1" in
    start) start_daemon ;;
    stop) stop_daemon ;;
    status)
        if [ -f "$PID_FILE" ] && ps -p $(cat "$PID_FILE") >/dev/null 2>&1; then
            echo "监控运行中 (PID: $(cat "$PID_FILE"))"
            [ -f "$LOG_FILE" ] && echo "最新日志:" && tail -1 "$LOG_FILE" 2>/dev/null || true
        else
            echo "监控未运行"
        fi
        ;;
    *) echo "使用: $0 {start|stop|status}" ;;
esac
EOF

    chmod +x "$SCRIPT_PATH"
    echo -e "${GREEN}✓ 脚本已安装到: $SCRIPT_PATH${NC}"
    echo -e "${YELLOW}使用方法:${NC}"
    echo -e "  ${CYAN}command_monitor.sh start${NC}   # 启动监控"
    echo -e "  ${CYAN}command_monitor.sh stop${NC}    # 停止监控"  
    echo -e "  ${CYAN}command_monitor.sh status${NC}  # 查看状态"
    return 0
}

# 设置开机自启动
setup_auto_start() {
    if [ ! -f "$SCRIPT_PATH" ]; then
        install_script
    fi
    
    # 使用crontab实现自启动（更可靠）
    (crontab -l 2>/dev/null | grep -v "command_monitor.sh"; echo "@reboot $SCRIPT_PATH start > /dev/null 2>&1") | crontab -
    
    echo -e "${GREEN}✓ 已设置开机自启动${NC}"
    echo -e "${YELLOW}重启后会自动启动监控服务${NC}"
}

# 取消开机自启动
remove_auto_start() {
    crontab -l 2>/dev/null | grep -v "command_monitor.sh" | crontab -
    echo -e "${YELLOW}✓ 已取消开机自动启动${NC}"
}

# 后台运行监控
start_background_monitor() {
    echo -e "${GREEN}启动后台监控服务...${NC}"
    
    # 如果已安装系统脚本，使用系统脚本
    if [ -f "$SCRIPT_PATH" ]; then
        $SCRIPT_PATH start
        return $?
    fi
    
    # 否则使用内置功能
    if [ -f "/tmp/command_monitor.pid" ]; then
        local old_pid=$(cat "/tmp/command_monitor.pid" 2>/dev/null)
        if ps -p "$old_pid" >/dev/null 2>&1; then
            echo -e "${YELLOW}监控服务已经在运行 (PID: $old_pid)${NC}"
            return 1
        fi
    fi
    
    (
        DAEMON_MODE=true
        init_log_system
        configure_realtime_history
        
        log_message "INFO" "后台监控进程启动 - PID: $$"
        log_message "INFO" "客户端IP: $(get_client_ip)"
        log_message "INFO" "地理位置: $(get_ip_location $(get_client_ip))"
        
        echo $$ > "/tmp/command_monitor.pid"
        monitor_history_files
    ) &
    
    local main_pid=$!
    sleep 2
    
    if ps -p $main_pid >/dev/null 2>&1; then
        echo -e "${GREEN}✓ 后台监控已启动!${NC}"
        echo -e "主进程PID: $main_pid"
        echo -e "日志文件: $LATEST_LOG"
        echo -e "查看日志: ${GREEN}tail -f $LATEST_LOG${NC}"
        return 0
    else
        echo -e "${RED}✗ 后台监控启动失败${NC}"
        return 1
    fi
}

# 前台运行监控
start_foreground_monitor() {
    echo -e "${YELLOW}前台运行模式 (输入 TO 切换后台，Ctrl+C 停止)${NC}"
    
    DAEMON_MODE=false
    init_log_system
    configure_realtime_history
    
    local client_ip=$(get_client_ip)
    local location_info=$(get_ip_location "$client_ip")
    
    log_message "INFO" "前台监控启动 - PID: $$"
    log_message "INFO" "客户端IP: $client_ip"
    log_message "INFO" "地理位置: $location_info"
    
    echo -e "${GREEN}开始监控...${NC}"
    echo -e "客户端IP: $client_ip"
    echo -e "地理位置: $location_info"
    echo -e "输入 ${CYAN}TO${NC} 切换到后台模式"
    
    monitor_history_files &
    local monitor_pid=$!
    
    while true; do
        if read -t 1 user_input; then
            if [ "$user_input" = "TO" ] || [ "$user_input" = "to" ]; then
                echo -e "\n${GREEN}切换到后台模式...${NC}"
                kill $monitor_pid 2>/dev/null || true
                start_background_monitor
                exit 0
            fi
        fi
        
        if ! ps -p $monitor_pid >/dev/null 2>&1; then
            echo -e "${RED}监控进程异常退出${NC}"
            break
        fi
    done
}

# 查看监控状态
check_monitor_status() {
    # 优先使用系统脚本检查
    if [ -f "$SCRIPT_PATH" ]; then
        $SCRIPT_PATH status
        return $?
    fi
    
    if [ -f "/tmp/command_monitor.pid" ]; then
        local main_pid=$(cat "/tmp/command_monitor.pid")
        if ps -p "$main_pid" >/dev/null 2>&1; then
            echo -e "${GREEN}✓ 监控服务运行中${NC}"
            echo "主进程PID: $main_pid"
            [ -f "$LATEST_LOG" ] && echo "最新日志:" && tail -1 "$LATEST_LOG" 2>/dev/null || true
            return 0
        fi
    fi
    
    echo -e "${RED}✗ 监控服务未运行${NC}"
    return 1
}

# 停止监控进程
stop_monitor() {
    # 优先使用系统脚本停止
    if [ -f "$SCRIPT_PATH" ]; then
        $SCRIPT_PATH stop
        return $?
    fi
    
    if [ -f "/tmp/command_monitor.pid" ]; then
        local main_pid=$(cat "/tmp/command_monitor.pid")
        echo -e "${YELLOW}正在停止监控进程 (PID: $main_pid)...${NC}"
        kill "$main_pid" 2>/dev/null || true
        sleep 1
        rm -f "/tmp/command_monitor.pid"
        echo -e "${GREEN}✓ 监控进程已停止${NC}"
    else
        pkill -f "monitor_history_files" 2>/dev/null && echo -e "${GREEN}✓ 监控进程已停止${NC}" || echo -e "${YELLOW}没有找到运行的监控进程${NC}"
    fi
}

# 显示交互式菜单
show_interactive_menu() {
    clear
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}       实时命令监控系统 v4.3${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo
    echo -e "${YELLOW}请选择操作:${NC}"
    echo
    echo -e "  ${CYAN}1${NC}. 启动后台监控模式"
    echo -e "  ${CYAN}2${NC}. 启动前台监控模式"
    echo -e "  ${CYAN}3${NC}. 查看监控状态"
    echo -e "  ${CYAN}4${NC}. 停止监控服务"
    echo -e "  ${CYAN}5${NC}. 设置开机自启动"
    echo -e "  ${CYAN}6${NC}. 取消开机自启动"
    echo -e "  ${CYAN}7${NC}. 安装脚本到系统"
    echo -e "  ${CYAN}8${NC}. 清理所有残留"
    echo -e "  ${CYAN}0${NC}. 退出"
    echo
    echo -e "${GREEN}========================================${NC}"
    echo -e "${YELLOW}提示:${NC}"
    echo -e "  - 前台模式下输入 ${CYAN}TO${NC} 可切换到后台"
    echo -e "  - 建议先选择8清理残留，再选择7安装"
    echo -e "${GREEN}========================================${NC}"
    echo
}

# 处理用户选择
handle_user_choice() {
    local choice
    read -p "请输入选择 [0-8]: " choice
    
    case $choice in
        1) start_background_monitor ;;
        2) start_foreground_monitor ;;
        3) check_monitor_status ;;
        4) stop_monitor ;;
        5) setup_auto_start ;;
        6) remove_auto_start ;;
        7) install_script ;;
        8) cleanup_all ;;
        0) echo -e "${GREEN}再见!${NC}"; exit 0 ;;
        *) echo -e "${RED}无效选择，请重新输入${NC}" ;;
    esac
    
    echo
    read -p "按回车键返回菜单..."
}

### 主程序 ###
main() {
    local command="${1:-}"
    
    case "$command" in
        -d|--daemon) start_background_monitor ;;
        -s|--status) check_monitor_status ;;
        -k|--kill) stop_monitor ;;
        -c|--clean) cleanup_all ;;
        -h|--help|"") show_interactive_menu; handle_user_choice ;;
        *) echo -e "${RED}未知选项: $command${NC}"; exit 1 ;;
    esac
}

# 启动
if [ $# -eq 0 ]; then
    while true; do
        show_interactive_menu
        handle_user_choice
    done
else
    main "$@"
fi