package main

import (
	"crypto/sha512"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io/fs"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
	"unsafe"

	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/sys/unix"
)

const (
	configFile     = "/etc/execguard/config.json"
	dbFile         = "/etc/execguard/allowed.db"
	logFile        = "/var/log/execguard.log"
	scanInterval   = 9 * time.Minute
)

type Config struct {
	ProtectedDirs []string `json:"protected_dirs"`
	AlertEmail    string   `json:"alert_email"`
    SkipDirs      []string `json:"skip_dirs"`
}

var initMode bool
var updateFile string
var config *Config

func main() {
	flag.BoolVar(&initMode, "init", false, "initialize and populate allowed executable database")
	flag.StringVar(&updateFile, "update", "", "add specified file to allowed database with hash")
	flag.Parse()

	if os.Geteuid() != 0 {
		log.Fatal("This program must be run as root")
        os.Exit(1) // Exit with status code 1
	}

	logf, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		log.Fatalf("Error opening log file: %v", err)
	}
	defer logf.Close()
	log.SetOutput(logf)

	db, err := sql.Open("sqlite3", dbFile)
	if err != nil {
		log.Fatalf("Error opening database: %v", err)
        os.Exit(2) // Exit with status code 2
	}
	defer db.Close()

	createTable(db)

	if updateFile != "" {
		absPath, err := filepath.Abs(updateFile)
		if err != nil {
			log.Fatalf("Invalid update file path: %v", err)
            os.Exit(1) // Exit with status code 1
		}
		addToAllowed(db, absPath)
		log.Printf("Added to allowed list: %s", absPath)
		return
	}

	config, err = loadConfig()
	if err != nil {
		log.Fatalf("Error loading config: %v", err)
        os.Exit(3) // Exit with status code 3
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from scan panic: %v", r)
			}
		}()
		periodicScan(config.ProtectedDirs, db)
	}()

	if err := monitorExecutions(db); err != nil {
		log.Fatalf("Execution monitoring failed: %v", err)
        os.Exit(4) // Exit with status code 4
	}
}

func loadConfig() (*Config, error) {
	data, err := os.ReadFile(configFile)
	if err != nil {
		return nil, err
	}
	var cfg Config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func createTable(db *sql.DB) {
	query := `CREATE TABLE IF NOT EXISTS allowed (
		path TEXT PRIMARY KEY,
		hash TEXT
	)`
	_, err := db.Exec(query)
	if err != nil {
		log.Fatalf("Failed to create table: %v", err)
        os.Exit(5) // Exit with status code 5
	}
}

func isAllowed(db *sql.DB, path string) bool {
	var storedHash string
	hash := computeHash(path)
	if hash == "" {
		return false
	}
	err := db.QueryRow("SELECT hash FROM allowed WHERE path = ?", path).Scan(&storedHash)
	return err == nil && storedHash == hash
}

func addToAllowed(db *sql.DB, path string) {
	hash := ""
	if initMode || updateFile != "" {
		hash = computeHash(path)
	}
	_, err := db.Exec("INSERT OR REPLACE INTO allowed(path, hash) VALUES(?, ?)", path, hash)
	if err != nil {
		log.Printf("Error inserting allowed entry: %v", err)
	}
}

func computeHash(path string) string {
	data, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	sha := sha512.Sum512(data)
	return hex.EncodeToString(sha[:])
}

func periodicScan(dirs []string, db *sql.DB) {
	skipSet := make(map[string]struct{})
	for _, skip := range config.SkipDirs {
		if abs, err := filepath.Abs(skip); err == nil {
			skipSet[abs] = struct{}{}
		}
	}

	for {
		for _, dir := range dirs {
			filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
				if err != nil {
					return nil
				}

				absPath, err := filepath.Abs(path)
				if err != nil {
					return nil
				}

				// Skip if in any of the SkipDirs
				for skipDir := range skipSet {
					if strings.HasPrefix(absPath, skipDir) {
						return filepath.SkipDir
					}
				}

				if d.Type().IsRegular() {
					info, err := d.Info()
					if err != nil || (info.Mode().Perm()&0111 == 0) {
						return nil
					}
					absPath, _ = filepath.EvalSymlinks(absPath)

					if initMode {
						addToAllowed(db, absPath)
					} else if !isAllowed(db, absPath) {
						log.Printf("Found unauthorized executable: %s", absPath)
						os.Chmod(absPath, info.Mode()&^0111)
						sendAlert(fmt.Sprintf("Unauthorized executable found and blocked: %s", absPath))
					}
				}
				return nil
			})
		}
		time.Sleep(scanInterval)
	}
}

func monitorExecutions(db *sql.DB) error {
	fd, err := unix.FanotifyInit(unix.FAN_CLOEXEC|unix.FAN_CLASS_CONTENT, unix.O_RDONLY|unix.O_LARGEFILE)
	if err != nil {
		return err
	}
	defer unix.Close(fd)

	for _, dir := range config.ProtectedDirs {
		if err := unix.FanotifyMark(fd, unix.FAN_MARK_ADD|unix.FAN_MARK_MOUNT,
			unix.FAN_OPEN_EXEC_PERM, unix.AT_FDCWD, dir); err != nil {
			log.Printf("Failed to mark %s: %v", dir, err)
		}
	}

	buf := make([]byte, 4096)
	for {
		n, err := unix.Read(fd, buf)
		if err != nil {
			return err
		}
		for offset := 0; offset < n; {
			meta := (*unix.FanotifyEventMetadata)(unsafe.Pointer(&buf[offset]))
			if meta.Event_len == 0 {
				break
			}

			resp := unix.FanotifyResponse{Fd: meta.Fd, Response: unix.FAN_ALLOW}
			defer unix.Close(int(meta.Fd))

			if meta.Mask&unix.FAN_OPEN_EXEC_PERM != 0 {
				fdpath := fmt.Sprintf("/proc/self/fd/%d", meta.Fd)
				path, err := os.Readlink(fdpath)
				if err == nil {
					absPath, _ := filepath.Abs(path)
					absPath, _ = filepath.EvalSymlinks(absPath)
					info, statErr := os.Stat(absPath)
					if statErr == nil && info.Mode().IsRegular() && (info.Mode().Perm()&0111 != 0) {
						if initMode {
							addToAllowed(db, absPath)
						} else if !isAllowed(db, absPath) {
							log.Printf("Blocked execution attempt: %s", absPath)
							sendAlert(fmt.Sprintf("Unauthorized execution attempt blocked: %s", absPath))
							resp.Response = unix.FAN_DENY
						}
					}
				}
			}

			b := (*[unsafe.Sizeof(resp)]byte)(unsafe.Pointer(&resp))[:]
			if _, err := unix.Write(fd, b); err != nil {
				log.Printf("Fanotify response write error: %v", err)
			}

			offset += int(meta.Event_len)
		}
	}
}

func sendAlert(message string) {
	if config.AlertEmail == "" {
		return
	}
	if _, err := exec.LookPath("mail"); err != nil {
		log.Printf("Mail command not found: %v", err)
		return
	}
	cmd := exec.Command("mail", "-s", "ExecGuard Alert", config.AlertEmail)
	cmd.Stdin = strings.NewReader(message)
	if err := cmd.Run(); err != nil {
		log.Printf("Failed to send alert: %v", err)
	}
}