package main

import (
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

type Credentials struct {
	AccessKey string `json:"access_key"`
	SecretKey string `json:"secret_key"`
	KeyArn    string `json:"key_arn"`
	Region    string `json:"region"`
}

func main() {
	if len(os.Args) < 2 {
		fmt.Printf("Usage: %s <keyserver> [socket]\n", os.Args[0])
		fmt.Printf("Example: %s https://keyserver.example.com\n", os.Args[0])
		fmt.Printf("Example: %s https://keyserver.example.com /var/run/kmsplugin/socket.sock\n", os.Args[0])
		os.Exit(1)
	}

	keyserver := os.Args[1]
	socket := "/var/run/kmsplugin/socket.sock"
	if len(os.Args) > 2 {
		socket = os.Args[2]
	}

	// Check if keyserver is reachable
	resp, err := http.Head(keyserver)
	if err != nil || resp.StatusCode >= 400 {
		log.Println("Key server not reachable, exiting.")
		os.Exit(1)
	}

	// Generate UUID
	uuid := getCPUModel() + getEth0MAC() + getSystemUUID()
	keyID := md5sum(uuid)

	log.Println("Downloading key from key server")
	keyURL := fmt.Sprintf("%s/%s-encryption-service-credentials.json", keyserver, keyID)
	resp, err = http.Get(keyURL)
	if err != nil || resp.StatusCode != 200 {
		log.Println("Failed to download key from key server")
		os.Exit(1)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body) // Changed from ioutil.ReadAll
	if err != nil {
		log.Println("Failed to read key file")
		os.Exit(1)
	}

	var creds Credentials
	if err := json.Unmarshal(body, &creds); err != nil {
		log.Println("Failed to extract credentials from key file")
		os.Exit(1)
	}

	os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKey)
	os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretKey)

	// Prepare socket dir
	socketDir := filepath.Dir(socket)
	os.MkdirAll(socketDir, 0700)
	os.Chmod(socketDir, 0700)

	if _, err := os.Stat("/usr/local/bin/aws-encryption-provider"); os.IsNotExist(err) {
		log.Println("Encryption provider not found, exiting.")
		os.Exit(1)
	}

	cmd := exec.Command("/usr/local/bin/aws-encryption-provider",
		"--key="+creds.KeyArn,
		"--region="+creds.Region,
		"--listen="+socket,
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		log.Println("Failed to start encryption provider")
		os.Exit(1)
	}
	os.Exit(0)
}

func getCPUModel() string {
	data, err := os.ReadFile("/proc/cpuinfo") // Changed from ioutil.ReadFile
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasPrefix(line, "model name") {
			return line
		}
	}
	return ""
}

func getEth0MAC() string {
	iface, err := net.InterfaceByName("eth0")
	if err != nil {
		return ""
	}
	return iface.HardwareAddr.String()
}

func getSystemUUID() string {
	out, err := exec.Command("sudo", "dmidecode", "-s", "system-uuid").Output()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out))
}

func md5sum(s string) string {
	h := md5.New()
	io.WriteString(h, s)
	return hex.EncodeToString(h.Sum(nil))
}