#!/usr/bin/env bash
# Advanced Bug Bounty Screenshot with Pre-made Annotation Templates
# Usage: bb-screenshot-annotate.sh <program> <type>
# Types: vulnerability, proof, request, response, comparison

set -euo pipefail

PROGRAM="${1:-current}"
TYPE="${2:-vulnerability}"
SCREENSHOT_DIR="${HOME}/bug-bounty/${PROGRAM}/screenshots"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
TEMP_FILE="/tmp/flameshot_${TIMESTAMP}.png"
FINAL_FILE="${SCREENSHOT_DIR}/${TIMESTAMP}_${TYPE}.png"

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

# Ensure screenshot directory exists
mkdir -p "$SCREENSHOT_DIR"

show_usage() {
    cat << EOF
Bug Bounty Screenshot Annotation Tool

Usage: bb-screenshot-annotate.sh <program> <type>

Screenshot Types:
  vulnerability   - Highlighting the vulnerability (red arrows/boxes)
  proof          - Proof of exploitation (green success indicators)
  request        - HTTP request in Burp Suite
  response       - HTTP response showing vulnerability
  comparison     - Before/After comparison
  evidence       - General evidence screenshot

Examples:
  bb-screenshot-annotate.sh juice-shop vulnerability
  bb-screenshot-annotate.sh acme proof
  bb-screenshot-annotate.sh target request

Tips:
  - Use Flameshot's built-in tools for annotation:
    * Arrow (for pointing)
    * Rectangle (for highlighting)
    * Text (for labels)
    * Pixelate (for redacting sensitive data)
  - Red for vulnerabilities
  - Green for successful exploitation
  - Yellow for important notes

EOF
}

if [[ "${1:-}" == "-h" ]] || [[ "${1:-}" == "--help" ]]; then
    show_usage
    exit 0
fi

# Display tips based on screenshot type
case "$TYPE" in
    vulnerability)
        echo -e "${RED}[!] VULNERABILITY SCREENSHOT${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Use RED arrows to point at the vulnerability"
        echo "  - Use RED rectangles to highlight affected areas"
        echo "  - Add text labels explaining what's wrong"
        ;;
    proof)
        echo -e "${GREEN}[!] PROOF OF EXPLOITATION${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Show successful exploitation result"
        echo "  - Highlight important output (session cookies, data, etc.)"
        echo "  - Use GREEN to show success"
        ;;
    request)
        echo -e "${BLUE}[!] HTTP REQUEST SCREENSHOT${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Capture Burp Suite request"
        echo "  - Highlight malicious payload in RED"
        echo "  - Show request method and endpoint clearly"
        ;;
    response)
        echo -e "${BLUE}[!] HTTP RESPONSE SCREENSHOT${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Capture server response"
        echo "  - Highlight vulnerability indicators (errors, data leaks)"
        echo "  - Show status code and response headers"
        ;;
    comparison)
        echo -e "${YELLOW}[!] BEFORE/AFTER COMPARISON${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Show side-by-side comparison"
        echo "  - Label 'BEFORE' and 'AFTER' clearly"
        echo "  - Highlight the difference"
        ;;
    evidence)
        echo -e "${BLUE}[!] GENERAL EVIDENCE${NC}"
        echo -e "${YELLOW}Tips:${NC}"
        echo "  - Capture relevant evidence"
        echo "  - Annotate important details"
        echo "  - Keep it clear and professional"
        ;;
    *)
        echo -e "${RED}[-] Unknown type: $TYPE${NC}"
        echo "Valid types: vulnerability, proof, request, response, comparison, evidence"
        exit 1
        ;;
esac

echo ""
echo -e "${BLUE}[+] Opening Flameshot...${NC}"
echo -e "${YELLOW}[!] Annotate your screenshot, then click Save${NC}"
echo ""

# Take screenshot with Flameshot (with annotation tools)
flameshot gui -p "$TEMP_FILE"

# Check if screenshot was actually taken
if [[ ! -f "$TEMP_FILE" ]]; then
    echo -e "${YELLOW}[!] Screenshot cancelled${NC}"
    exit 0
fi

echo -e "${GREEN}[✓] Screenshot captured${NC}"
echo -e "${BLUE}[+] Adding professional border and shadow...${NC}"

# Add border based on type
case "$TYPE" in
    vulnerability)
        BORDER_COLOR='#DC143C'  # Crimson red
        ;;
    proof)
        BORDER_COLOR='#228B22'  # Forest green
        ;;
    request|response)
        BORDER_COLOR='#4169E1'  # Royal blue
        ;;
    comparison)
        BORDER_COLOR='#FF8C00'  # Dark orange
        ;;
    evidence)
        BORDER_COLOR='#696969'  # Dim gray
        ;;
esac

# Add colored border, white mat, outer border, and drop shadow
convert "$TEMP_FILE" \
    -bordercolor "$BORDER_COLOR" -border 3 \
    -bordercolor white -border 12 \
    -bordercolor '#333333' -border 1 \
    \( +clone -background black -shadow 80x5+8+8 \) \
    +swap -background white -layers merge +repage \
    "$FINAL_FILE"

# Remove temp file
rm "$TEMP_FILE"

echo -e "${GREEN}[✓] Screenshot saved: $FINAL_FILE${NC}"
echo -e "${BLUE}[+] Copying path to clipboard...${NC}"

# Copy relative path to clipboard
RELATIVE_PATH="screenshots/$(basename "$FINAL_FILE")"
echo "$RELATIVE_PATH" | xclip -selection clipboard 2>/dev/null || \
echo "$RELATIVE_PATH" | wl-copy 2>/dev/null || \
echo -e "${YELLOW}[!] Could not copy to clipboard${NC}"

echo ""
echo -e "${GREEN}Path: $RELATIVE_PATH${NC}"
echo -e "${YELLOW}[!] Paste this into your vulnerability JSON:${NC}"
echo -e '    "path": "'"$RELATIVE_PATH"'",'

# Show in file manager
if command -v xdg-open &> /dev/null; then
    xdg-open "$(dirname "$FINAL_FILE")" &
fi