package policy

import (
	"go/parser"
	"go/token"
	"io/fs"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

func TestExecImportsStayInsideApprovedPackages(t *testing.T) {
	t.Parallel()

	_, thisFile, _, ok := runtime.Caller(0)
	if !ok {
		t.Fatal("runtime.Caller failed")
	}
	repoRoot := filepath.Clean(filepath.Join(filepath.Dir(thisFile), "..", ".."))
	fset := token.NewFileSet()

	var offenders []string
	err := filepath.WalkDir(filepath.Join(repoRoot, "internal"), func(path string, entry fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if entry.IsDir() {
			return nil
		}
		if filepath.Ext(path) != ".go" || strings.HasSuffix(path, "_test.go") {
			return nil
		}
		relPath, err := filepath.Rel(repoRoot, path)
		if err != nil {
			return err
		}
		if allowedExecImportPath(relPath) {
			return nil
		}
		file, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly)
		if err != nil {
			return err
		}
		for _, imp := range file.Imports {
			if imp.Path != nil && imp.Path.Value == `"os/exec"` {
				offenders = append(offenders, relPath)
				break
			}
		}
		return nil
	})
	if err != nil {
		t.Fatalf("walk repo: %v", err)
	}
	if len(offenders) != 0 {
		t.Fatalf("os/exec imports are only allowed in internal/cli, internal/firecracker, internal/system, and internal/vsockagent; found %v", offenders)
	}
}

func allowedExecImportPath(relPath string) bool {
	return strings.HasPrefix(relPath, "internal/cli/") ||
		strings.HasPrefix(relPath, "internal/firecracker/") ||
		strings.HasPrefix(relPath, "internal/system/") ||
		strings.HasPrefix(relPath, "internal/vsockagent/")
}