#!/bin/bash

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

install_nmap() {
    if ! command -v nmap &> /dev/null; then
        echo -e "${YELLOW}nmap could not be found, attempting to install...${NC}"
        if [[ "$OSTYPE" == "linux-gnu"* ]]; then
            if command -v apt &> /dev/null; then
                sudo apt update && sudo apt install -y nmap
            elif command -v yum &> /dev/null; then
                sudo yum install -y nmap
            elif command -v dnf &> /dev/null; then
                sudo dnf install -y nmap
            else
                echo -e "${RED}Unsupported package manager. Please install nmap manually.${NC}"
                exit 1
            fi
        elif [[ "$OSTYPE" == "darwin"* ]]; then
            if command -v brew &> /dev/null; then
                brew install nmap
            else
                echo -e "${RED}Homebrew is not installed. Please install Homebrew and then nmap.${NC}"
                exit 1
            fi
        else
            echo -e "${RED}Unsupported OS. Please install nmap manually.${NC}"
            exit 1
        fi
    fi
}

check_local_openssh_version() {
    local_version=$(ssh -V 2>&1 | awk '{print $1}')
    echo -e "${BLUE}Local machine OpenSSH version: $local_version${NC}"

    vulnerable_versions=(
        'OpenSSH_8.5'
        'OpenSSH_8.6'
        'OpenSSH_8.7'
        'OpenSSH_8.8'
        'OpenSSH_8.9'
        'OpenSSH_9.0'
        'OpenSSH_9.1'
        'OpenSSH_9.2'
        'OpenSSH_9.3'
        'OpenSSH_9.4'
        'OpenSSH_9.5'
        'OpenSSH_9.6'
        'OpenSSH_9.7'
    )

    for version in "${vulnerable_versions[@]}"; do
        if [[ "$local_version" == "$version" ]]; then
            echo -e "${RED}Local machine is running a vulnerable version of OpenSSH: $local_version${NC}"
            return
        fi
    done

    echo -e "${GREEN}Local machine is not running a vulnerable version of OpenSSH.${NC}"
}

get_ssh_banner() {
    local ip=$1
    local port=$2
    banner=$(nmap -Pn -p $port --script=banner $ip 2>/dev/null | grep -i "banner:")
    echo "$banner"
}

check_vulnerability() {
    local ip=$1
    local port=$2
    local banner=$(get_ssh_banner "$ip" "$port")
    if [[ -z "$banner" ]]; then
        echo -e "${YELLOW}$ip:$port:closed:Port closed${NC}"
        return
    fi

    if ! echo "$banner" | grep -q "SSH-2.0-OpenSSH"; then
        echo -e "${RED}$ip:$port:failed:Failed to retrieve SSH banner: $banner${NC}"
        return
    fi

    vulnerable_versions=(
        'SSH-2.0-OpenSSH_8.5'
        'SSH-2.0-OpenSSH_8.6'
        'SSH-2.0-OpenSSH_8.7'
        'SSH-2.0-OpenSSH_8.8'
        'SSH-2.0-OpenSSH_8.9'
        'SSH-2.0-OpenSSH_9.0'
        'SSH-2.0-OpenSSH_9.1'
        'SSH-2.0-OpenSSH_9.2'
        'SSH-2.0-OpenSSH_9.3'
        'SSH-2.0-OpenSSH_9.4'
        'SSH-2.0-OpenSSH_9.5'
        'SSH-2.0-OpenSSH_9.6'
        'SSH-2.0-OpenSSH_9.7'
    )

    excluded_versions=(
        'SSH-2.0-OpenSSH_8.9p1 Ubuntu-3ubuntu0.10'
        'SSH-2.0-OpenSSH_9.3p1 Ubuntu-3ubuntu3.6'
        'SSH-2.0-OpenSSH_9.6p1 Ubuntu-3ubuntu13.3'
        'SSH-2.0-OpenSSH_9.3p1 Ubuntu-1ubuntu3.6'
        'SSH-2.0-OpenSSH_9.2p1 Debian-2+deb12u3'
        'SSH-2.0-OpenSSH_8.4p1 Debian-5+deb11u3'
    )

    for version in "${vulnerable_versions[@]}"; do
        if echo "$banner" | grep -q "$version"; then
            for excluded in "${excluded_versions[@]}"; do
                if echo "$banner" | grep -q "$excluded"; then
                    echo -e "${GREEN}$ip:$port:not_vulnerable:(running $banner)${NC}"
                    return
                fi
            done
            echo -e "${RED}$ip:$port:vulnerable:(running $banner)${NC}"
            return
        fi
    done

    echo -e "${GREEN}$ip:$port:not_vulnerable:(running $banner)${NC}"
}

print_help() {
    echo -e "${BLUE}CVE-2024-6387 Checker${NC}"
    echo -e "${NC}Star on Github: https://github.com/devarshishimpi/CVE-2024-6387-Check${NC}"
    echo -e "${NC}${NC}"
    echo -e "${GREEN}Usage: ./check.sh [options] [target(s)]${NC}"
    echo -e "${NC}${NC}"
    echo -e "${YELLOW}Options:${NC}"
    echo -e "${NC}  -p, --port [PORT]       Specify the port to scan (default is 22)${NC}"
    echo -e "${NC}  -h, --help              Show this help message and exit${NC}"
    echo -e "${YELLOW}Target(s):${NC}"
    echo -e "${NC}  IP address(es) or hostname(s) of the target(s) to scan${NC}"
    echo -e "${YELLOW}Examples:${NC}"
    echo -e "${NC}  ./check.sh 192.168.1.1${NC}"
    echo -e "${NC}  ./check.sh -p 2222 192.168.1.1 192.168.1.2${NC}"
    echo -e "${NC}  ./check.sh --help${NC}"
}

main() {
    if [[ "$1" == "--help" || "$1" == "-h" || "$#" == 0 ]]; then
        print_help
        exit 0
    fi

    echo -e "${BLUE}Detected OS: $OSTYPE${NC}"
    
    install_nmap

    check_local_openssh_version

    local targets=()
    local port=22

    while [[ "$#" -gt 0 ]]; do
        case "$1" in
            --port|-p)
                port=$2
                shift 2
                ;;
            *)
                targets+=("$1")
                shift
                ;;
        esac
    done

    if [[ ${#targets[@]} -eq 0 ]]; then
        print_help
        exit 1
    fi

    closed_ports=0
    not_vulnerable=()
    vulnerable=()

    for target in "${targets[@]}"; do
        ip=$(echo "$target" | tr -d '[:space:]')
        result=$(check_vulnerability "$ip" "$port")
        IFS=':' read -r ip port status message <<< "$result"
        case "$status" in
            closed)
                closed_ports=$((closed_ports + 1))
                ;;
            vulnerable)
                vulnerable+=("$ip ($message)")
                ;;
            not_vulnerable)
                not_vulnerable+=("$ip ($message)")
                ;;
            *)
                echo -e "${YELLOW}⚠️ [!] Server at $ip:$port is $message${NC}"
                ;;
        esac
    done

    echo -e "\n${GREEN}Servers not vulnerable: ${#not_vulnerable[@]}${NC}\n"
    for item in "${not_vulnerable[@]}"; do
        echo -e "${GREEN}   [+] Server at $item${NC}"
    done

    echo -e "\n${RED}Servers likely vulnerable: ${#vulnerable[@]}${NC}\n"
    for item in "${vulnerable[@]}"; do
        echo -e "${RED}   [+] Server at $item${NC}"
    done

    echo -e "\n${YELLOW}Servers with port $port closed: $closed_ports${NC}\n"
}

main "$@"
