#!/bin/bash
#
# audited.sh Local Audit Tool
# https://audited.sh
#
# Run security audits locally using your Claude CLI subscription.
# For professional audits with badges, visit https://audited.sh
#

set -e

VERSION="1.6.0"

# TODO: Change to https://audited.sh for production
BASE_URL="http://localhost:3000"
SCRIPT_URL="$BASE_URL/full"
API_URL="$BASE_URL"

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

print_banner() {
    echo -e "${BLUE}"
    echo "  ╔═══════════════════════════════════════╗"
    echo "  ║     audited.sh Local Audit Tool      ║"
    echo "  ║           v${VERSION}                       ║"
    echo "  ╚═══════════════════════════════════════╝"
    echo -e "${NC}"
}

auto_update() {
    # Skip if already updated this run (prevents infinite loop)
    [ "$AUDITED_UPDATED" = "1" ] && return

    local tmp_file=$(mktemp)
    if ! curl -sL --connect-timeout 3 "$SCRIPT_URL" -o "$tmp_file" 2>/dev/null; then
        rm -f "$tmp_file"
        return
    fi

    local remote_version=$(grep '^VERSION=' "$tmp_file" | head -1 | cut -d'"' -f2)
    if [ -z "$remote_version" ] || [ "$remote_version" = "$VERSION" ]; then
        rm -f "$tmp_file"
        return
    fi

    echo -e "${YELLOW}Updating to v${remote_version}...${NC}"

    local script_path="$0"
    if [ -L "$script_path" ]; then
        script_path=$(readlink -f "$script_path")
    fi

    chmod +x "$tmp_file"
    mv "$tmp_file" "$script_path"

    echo -e "${GREEN}Updated! Restarting...${NC}"
    echo ""

    # Re-execute with same arguments
    AUDITED_UPDATED=1 exec "$script_path" "$@"
}

print_usage() {
    echo "Usage: $0 <--local|--api> [OPTIONS] [directory]"
    echo ""
    echo "Modes (one required):"
    echo "  --local              Run locally using Claude CLI (requires Pro/Max subscription)"
    echo "  --api                Use audited.sh API (requires AUDITED_API_KEY)"
    echo ""
    echo "Options:"
    echo "  -o, --output FILE    Output file (default: audit-report.md)"
    echo "  -h, --help           Show this help message"
    echo "  -v, --version        Show version"
    echo ""
    echo "Environment variables:"
    echo "  AUDITED_API_KEY      API key (required for --api mode)"
    echo "  AUDITED_CODE         Promo code (free) or referral code (10% off)"
    echo ""
    echo "Examples:"
    echo "  $0 --local           # Local audit of current directory"
    echo "  $0 --api ./src       # API audit of ./src"
    echo "  $0 --local -o report.md ./project"
}

# Run audit via audited.sh API
run_api_audit() {
    local dir="$1"
    local output="$2"

    echo -e "${YELLOW}Collecting files...${NC}"

    local files=$(collect_files "$dir")
    local file_count=$(echo "$files" | grep -c . || echo "0")

    if [ "$file_count" -eq 0 ]; then
        echo -e "${RED}Error: No source files found in $dir${NC}"
        exit 1
    fi

    echo -e "Found ${GREEN}$file_count${NC} files to analyze"

    # Create zip file
    echo -e "${YELLOW}Creating archive...${NC}"
    local zip_file=$(mktemp).zip

    # Change to directory and create zip with relative paths
    (cd "$dir" && echo "$files" | sed "s|^$dir/||" | sed "s|^\./||" | zip -q "$zip_file" -@)

    # Get project name from directory
    local project_name=$(basename "$(cd "$dir" && pwd)")

    # Submit to API
    echo -e "${YELLOW}Submitting to audited.sh API...${NC}"
    local curl_args=(-s -X POST "$API_URL/api/v1/audits" \
        -H "Authorization: Bearer $AUDITED_API_KEY" \
        -F "file=@$zip_file" \
        -F "name=$project_name")

    # Add code if provided via environment variable
    if [ -n "$AUDITED_CODE" ]; then
        curl_args+=(-F "code=$AUDITED_CODE")
    fi

    local response=$(curl "${curl_args[@]}")

    rm -f "$zip_file"

    local audit_id=$(echo "$response" | grep -o '"id":[0-9]*' | head -1 | cut -d: -f2)

    if [ -z "$audit_id" ]; then
        echo -e "${RED}Error: Failed to create audit${NC}"
        echo "$response"
        exit 1
    fi

    echo -e "Audit ID: ${GREEN}$audit_id${NC}"

    # Poll for completion
    echo -e "${YELLOW}Waiting for audit to complete...${NC}"
    local audit_status="pending"
    local attempts=0
    local max_attempts=120  # 10 minutes max

    while [ "$audit_status" != "completed" ] && [ "$audit_status" != "failed" ] && [ $attempts -lt $max_attempts ]; do
        sleep 5
        attempts=$((attempts + 1))

        local status_response=$(curl -s "$API_URL/api/v1/audits/$audit_id" \
            -H "Authorization: Bearer $AUDITED_API_KEY")

        audit_status=$(echo "$status_response" | grep -o '"status":"[^"]*"' | head -1 | cut -d'"' -f4)
        echo -ne "\r  Status: $audit_status (${attempts}/${max_attempts})  "
    done
    echo ""

    if [ "$audit_status" = "failed" ]; then
        echo -e "${RED}Error: Audit failed${NC}"
        exit 1
    fi

    if [ "$audit_status" != "completed" ]; then
        echo -e "${RED}Error: Audit timed out${NC}"
        exit 1
    fi

    # Get report
    echo -e "${YELLOW}Fetching report...${NC}"
    curl -s "$API_URL/api/v1/audits/$audit_id/report" \
        -H "Authorization: Bearer $AUDITED_API_KEY" > "$output"

    echo -e "${GREEN}✓ Audit complete!${NC}"
    echo -e "Report saved to: ${BLUE}$output${NC}"
    echo ""

    # Get findings summary
    local findings_response=$(curl -s "$API_URL/api/v1/audits/$audit_id/findings" \
        -H "Authorization: Bearer $AUDITED_API_KEY")

    echo -e "${YELLOW}=== FINDINGS SUMMARY ===${NC}"
    echo ""

    local critical=$(echo "$findings_response" | grep -o '"critical":[0-9]*' | cut -d: -f2)
    local high=$(echo "$findings_response" | grep -o '"high":[0-9]*' | cut -d: -f2)
    local medium=$(echo "$findings_response" | grep -o '"medium":[0-9]*' | cut -d: -f2)
    local low=$(echo "$findings_response" | grep -o '"low":[0-9]*' | cut -d: -f2)
    local info=$(echo "$findings_response" | grep -o '"informational":[0-9]*' | cut -d: -f2)

    [ "${critical:-0}" -gt 0 ] && echo -e "${RED}Critical: $critical${NC}"
    [ "${high:-0}" -gt 0 ] && echo -e "${RED}High: $high${NC}"
    [ "${medium:-0}" -gt 0 ] && echo -e "${YELLOW}Medium: $medium${NC}"
    [ "${low:-0}" -gt 0 ] && echo -e "${BLUE}Low: $low${NC}"
    [ "${info:-0}" -gt 0 ] && echo -e "Informational: $info"

    echo ""
    echo -e "View online: ${BLUE}$API_URL/audits/$audit_id${NC}"
}

check_claude_cli() {
    if ! command -v claude &> /dev/null; then
        echo -e "${RED}Error: Claude CLI not found${NC}"
        echo ""
        echo "Please install Claude CLI from: https://claude.ai/download"
        echo "Make sure you have an active Claude Pro or Max subscription."
        exit 1
    fi
}

collect_files() {
    local dir="$1"

    # Collect all supported source files
    find "$dir" -type f \( \
        -name "*.sol" \
        -o -name "*.rb" \
        -o -name "*.erb" \
        -o -name "*.haml" \
        -o -name "*.py" \
        -o -name "*.js" \
        -o -name "*.ts" \
        -o -name "*.jsx" \
        -o -name "*.tsx" \
        -o -name "*.go" \
        -o -name "*.rs" \
        -o -name "*.java" \
        -o -name "*.kt" \
        -o -name "*.swift" \
        -o -name "*.c" \
        -o -name "*.cpp" \
        -o -name "*.h" \
    \) \
        ! -path "*/node_modules/*" \
        ! -path "*/vendor/*" \
        ! -path "*/.bundle/*" \
        ! -path "*/.venv/*" \
        ! -path "*/venv/*" \
        ! -path "*/__pycache__/*" \
        ! -path "*/dist/*" \
        ! -path "*/build/*" \
        ! -path "*/lib/*" \
        ! -path "*/.git/*" \
        ! -name "*.min.js" \
        | sort
}

get_prompt() {
    cat <<'PROMPT'
You are a security auditor. Analyze the provided code for security vulnerabilities.

OUTPUT FORMAT (Markdown):

# Security Audit Report

## Executive Summary
Brief overview of files reviewed and key findings.

## Findings Summary
| Severity | ID | Title | Location |
|----------|-----|-------|----------|
(list all findings)

## Findings

For each finding:

### [SEVERITY-#] Title

**Severity:** Critical/High/Medium/Low/Informational

**Location:** `filename:line`

**Description:**
Explain the vulnerability.

**Impact:**
What could happen if exploited.

**Recommendation:**
How to fix it.

---

## Conclusion
Overall assessment and recommendations.

SEVERITY GUIDE:
- Critical: Direct exploitation possible, severe impact
- High: Significant risk, likely exploitable
- Medium: Moderate risk, requires conditions
- Low: Minor issues, best practices
- Informational: Suggestions, no direct risk

Analyze the code below and provide your security audit report:

PROMPT
}

run_audit() {
    local dir="$1"
    local output="$2"

    echo -e "${YELLOW}Collecting files...${NC}"

    local files=$(collect_files "$dir")
    local file_count=$(echo "$files" | grep -c . || echo "0")

    if [ "$file_count" -eq 0 ]; then
        echo -e "${RED}Error: No source files found in $dir${NC}"
        exit 1
    fi

    echo -e "Found ${GREEN}$file_count${NC} files to analyze"

    # Build the prompt with code
    local prompt=$(get_prompt)

    echo -e "${YELLOW}Building code bundle...${NC}"

    local code=""
    while IFS= read -r file; do
        if [ -f "$file" ]; then
            local content=$(cat "$file" 2>/dev/null | head -c 50000)  # Limit per file
            code+="=== $file ===
$content

"
        fi
    done <<< "$files"

    # Truncate if too long
    if [ ${#code} -gt 150000 ]; then
        code="${code:0:150000}
[truncated - code exceeds 150KB limit]"
    fi

    local full_prompt="$prompt

$code"

    echo -e "${YELLOW}Running audit with Claude CLI...${NC}"
    echo -e "${BLUE}This may take a few minutes depending on code size.${NC}"
    echo ""

    # Run Claude CLI
    echo "$full_prompt" | claude --print > "$output" 2>&1

    if [ $? -eq 0 ] && [ -s "$output" ]; then
        echo -e "${GREEN}✓ Audit complete!${NC}"
        echo -e "Report saved to: ${BLUE}$output${NC}"
        echo ""

        # Extract and display findings summary
        echo -e "${YELLOW}=== FINDINGS SUMMARY ===${NC}"
        echo ""

        # Count findings by severity
        local critical=0 high=0 medium=0 low=0 info=0
        critical=$(grep -c "| Critical |" "$output" 2>/dev/null) || true
        high=$(grep -c "| High |" "$output" 2>/dev/null) || true
        medium=$(grep -c "| Medium |" "$output" 2>/dev/null) || true
        low=$(grep -c "| Low |" "$output" 2>/dev/null) || true
        info=$(grep -c "| Informational |" "$output" 2>/dev/null) || true

        [ "$critical" -gt 0 ] && echo -e "${RED}Critical: $critical${NC}"
        [ "$high" -gt 0 ] && echo -e "${RED}High: $high${NC}"
        [ "$medium" -gt 0 ] && echo -e "${YELLOW}Medium: $medium${NC}"
        [ "$low" -gt 0 ] && echo -e "${BLUE}Low: $low${NC}"
        [ "$info" -gt 0 ] && echo -e "Informational: $info"

        echo ""

        # Extract findings summary table
        if grep -q "^| Severity |" "$output"; then
            echo -e "${YELLOW}Findings:${NC}"
            # Print the findings table (header + all finding rows)
            sed -n '/^| Severity |/,/^$/p' "$output" | head -20
            echo ""
        fi

        # Show next steps
        if [ "$critical" -gt 0 ] || [ "$high" -gt 0 ]; then
            echo -e "${RED}⚠ Critical/High severity issues found. Review and fix before deployment.${NC}"
        elif [ "$medium" -gt 0 ]; then
            echo -e "${YELLOW}Medium severity issues found. Consider addressing these.${NC}"
        else
            echo -e "${GREEN}No critical issues found.${NC}"
        fi

        echo ""
        echo -e "Full report: ${BLUE}$output${NC}"
        echo -e "For professional audits with public badges, visit ${BLUE}https://audited.sh${NC}"
    else
        echo -e "${RED}Error: Audit failed${NC}"
        cat "$output"
        exit 1
    fi
}

# Save original arguments for auto-update re-exec
ORIGINAL_ARGS=("$@")

# Parse arguments
OUTPUT="audit-report.md"
DIR=""
MODE=""

while [[ $# -gt 0 ]]; do
    case $1 in
        --local)
            MODE="local"
            shift
            ;;
        --api)
            MODE="api"
            shift
            ;;
        -o|--output)
            OUTPUT="$2"
            shift 2
            ;;
        -h|--help)
            print_usage
            exit 0
            ;;
        -v|--version)
            echo "audited.sh Local Audit Tool v${VERSION}"
            exit 0
            ;;
        -*)
            echo -e "${RED}Unknown option: $1${NC}"
            print_usage
            exit 1
            ;;
        *)
            DIR="$1"
            shift
            ;;
    esac
done

# Main
print_banner
auto_update "${ORIGINAL_ARGS[@]}"

# Require mode to be specified
if [ -z "$MODE" ]; then
    echo -e "${RED}Error: Please specify --local or --api${NC}"
    echo ""
    print_usage
    exit 1
fi

# Default to current directory
if [ -z "$DIR" ]; then
    DIR="."
fi

if [ ! -d "$DIR" ]; then
    echo -e "${RED}Error: Directory not found: $DIR${NC}"
    exit 1
fi

# Run the selected mode
if [ "$MODE" = "api" ]; then
    if [ -z "$AUDITED_API_KEY" ]; then
        echo -e "${RED}Error: AUDITED_API_KEY environment variable is required for API mode${NC}"
        echo ""
        echo "Get your API key at: https://audited.sh/api/keys"
        exit 1
    fi
    echo -e "${GREEN}Using audited.sh API${NC}"
    echo ""
    run_api_audit "$DIR" "$OUTPUT"
else
    check_claude_cli
    run_audit "$DIR" "$OUTPUT"
fi