move-data-nas/internal/api/handlers_sshkeys.go

package api

import (
	"crypto/sha256"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/go-chi/chi/v5"
	"github.com/syncserver/internal/config"
	"github.com/syncserver/internal/models"
	"github.com/syncserver/internal/sshmanager"
)

type SSHKeyHandler struct {
	db  *sql.DB
	cfg *config.Config
}

func NewSSHKeyHandler(db *sql.DB, cfg *config.Config) *SSHKeyHandler {
	return &SSHKeyHandler{db: db, cfg: cfg}
}

func (h *SSHKeyHandler) List(w http.ResponseWriter, r *http.Request) {
	repo := models.NewSSHKeyRepository(h.db)
	machineRepo := models.NewMachineRepository(h.db)
	keys, err := repo.GetAll()
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to fetch ssh keys")
		return
	}
	machines, _ := machineRepo.GetAll()

	out := make([]SSHKeyResponse, len(keys))
	for i, k := range keys {
		inUse := false
		for _, m := range machines {
			if m.SSHKeyID != nil && *m.SSHKeyID == k.ID {
				inUse = true
				break
			}
		}
		hasPriv := false
		if _, err := os.Stat(k.PrivateKeyPath); err == nil {
			hasPriv = true
		}
		fp, _ := sshmanager.Fingerprint(k.PublicKey)
		if fp == "" && hasPriv {
			pubKey, newFP, _ := sshmanager.RegeneratePublicKeyFromPrivateKeyFile(k.PrivateKeyPath, k.Label)
			repo.UpdatePublicKey(k.ID, pubKey)
			k.PublicKey = pubKey
			fp = newFP
		}
		out[i] = SSHKeyResponse{
			ID:            k.ID,
			Label:         k.Label,
			PublicKey:     k.PublicKey,
			Fingerprint:   fp,
			InUse:         inUse,
			HasPrivateKey: hasPriv,
			CreatedAt:     k.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
		}
	}
	writeJSON(w, out)
}

func (h *SSHKeyHandler) Create(w http.ResponseWriter, r *http.Request) {
	var req SSHKeyRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid request body")
		return
	}
	if req.Label == "" {
		writeError(w, http.StatusBadRequest, "label is required")
		return
	}
	if !req.Generate && req.PublicKey == "" {
		writeError(w, http.StatusBadRequest, "either generate=true or public_key is required")
		return
	}

	keysDir := filepath.Join(h.cfg.SSHDir(), "keys")
	if err := os.MkdirAll(keysDir, 0700); err != nil {
		writeError(w, http.StatusInternalServerError, "failed to create keys directory")
		return
	}

	var pubKey, privPath, fp string
	var err error

	if req.Generate {
		privPath, _, pubKey, fp, err = sshmanager.GenerateKeyPair(req.Label, keysDir)
		if err != nil {
			writeError(w, http.StatusInternalServerError, fmt.Sprintf("generating key: %v", err))
			return
		}
	} else {
		pubKey = strings.TrimSpace(req.PublicKey)
		fp, err = sshmanager.Fingerprint(pubKey)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid public key format")
			return
		}
		privPath = ""
	}

	repo := models.NewSSHKeyRepository(h.db)
	id, err := repo.Create(req.Label, privPath, pubKey)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to store ssh key")
		return
	}
	w.Header().Set("Location", "/api/ssh-keys/"+strconv.FormatInt(id, 10))
	writeJSON(w, SSHKeyResponse{
		ID:            id,
		Label:         req.Label,
		PublicKey:     pubKey,
		Fingerprint:   fp,
		InUse:         false,
		HasPrivateKey: privPath != "",
	}, http.StatusCreated)
}

func (h *SSHKeyHandler) Get(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid id")
		return
	}
	repo := models.NewSSHKeyRepository(h.db)
	k, err := repo.GetByID(id)
	if err == sql.ErrNoRows {
		writeError(w, http.StatusNotFound, "ssh key not found")
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to fetch ssh key")
		return
	}
	machineRepo := models.NewMachineRepository(h.db)
	machines, _ := machineRepo.GetAll()
	inUse := false
	for _, m := range machines {
		if m.SSHKeyID != nil && *m.SSHKeyID == k.ID {
			inUse = true
			break
		}
	}
	hasPriv := false
	if _, err := os.Stat(k.PrivateKeyPath); err == nil {
		hasPriv = true
	}
	fp, _ := sshmanager.Fingerprint(k.PublicKey)
	if fp == "" && hasPriv {
		pubKey, newFP, _ := sshmanager.RegeneratePublicKeyFromPrivateKeyFile(k.PrivateKeyPath, k.Label)
		repo.UpdatePublicKey(k.ID, pubKey)
		k.PublicKey = pubKey
		fp = newFP
	}
	writeJSON(w, SSHKeyResponse{
		ID:            k.ID,
		Label:         k.Label,
		PublicKey:     k.PublicKey,
		Fingerprint:   fp,
		InUse:         inUse,
		HasPrivateKey: hasPriv,
		CreatedAt:     k.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
	})
}

func (h *SSHKeyHandler) Delete(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid id")
		return
	}
	repo := models.NewSSHKeyRepository(h.db)
	machineRepo := models.NewMachineRepository(h.db)
	machines, _ := machineRepo.GetAll()
	for _, m := range machines {
		if m.SSHKeyID != nil && *m.SSHKeyID == id {
			writeError(w, http.StatusConflict, "ssh key is in use by machines")
			return
		}
	}
	k, err := repo.GetByID(id)
	if err == nil && k.PrivateKeyPath != "" {
		os.Remove(k.PrivateKeyPath)
		os.Remove(k.PrivateKeyPath + ".pub")
	}
	if err := repo.Delete(id); err != nil {
		writeError(w, http.StatusInternalServerError, "failed to delete ssh key")
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *SSHKeyHandler) DownloadPrivate(w http.ResponseWriter, r *http.Request) {
	id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid id")
		return
	}
	repo := models.NewSSHKeyRepository(h.db)
	k, err := repo.GetByID(id)
	if err == sql.ErrNoRows {
		writeError(w, http.StatusNotFound, "ssh key not found")
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to fetch ssh key")
		return
	}
	if k.PrivateKeyPath == "" {
		writeError(w, http.StatusNotFound, "no private key available for this entry")
		return
	}

	data, err := os.ReadFile(k.PrivateKeyPath)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to read private key")
		return
	}

	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s.key"`, k.Label))
	w.Header().Set("X-Private-Key-Hash", fmt.Sprintf("sha256:%x", sha256.Sum256(data)))
	io.WriteString(w, string(data))
}