package controllers

import (
	"crypto/rand"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"code.crute.us/mcrute/ssh-proxy/app/middleware"
	"code.crute.us/mcrute/ssh-proxy/app/models"
	"github.com/labstack/echo/v4"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"golang.org/x/crypto/ssh"
)

var (
	caError = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "ssh_proxy",
		Name:      "ssh_ca_error",
		Help:      "Total number of errors during SSH CA operation",
	}, []string{"type"})
	caSuccess = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "ssh_proxy",
		Name:      "ssh_ca_success",
		Help:      "Total number of successful CA operations",
	})
)

type CASecret struct {
	Key string `mapstructure:"key"`
}

type CAHandlerConfig struct {
	Logger     echo.Logger
	Users      models.UserStore
	Expiration time.Duration
	Secret     CASecret
}

type CAHandler struct {
	Logger     echo.Logger
	Users      models.UserStore
	Expiration time.Duration
	signer     ssh.Signer
}

func NewCAHandler(cfg CAHandlerConfig) (*CAHandler, error) {
	signer, err := ssh.ParsePrivateKey([]byte(cfg.Secret.Key))
	if err != nil {
		return nil, err
	}

	cfg.Logger.Infof("CA Authorized Key: %s", ssh.MarshalAuthorizedKey(signer.PublicKey()))

	return &CAHandler{
		Logger:     cfg.Logger,
		Users:      cfg.Users,
		Expiration: cfg.Expiration,
		signer:     signer,
	}, nil
}

func (h *CAHandler) authorizeRequest(c echo.Context, certRequest *ssh.Certificate) error {
	session := middleware.GetAuthorizedSession(c)

	user, err := h.Users.Get(c.Request().Context(), session.UserId)
	if err != nil {
		return err
	}

	if user.Username != certRequest.ValidPrincipals[0] {
		caError.With(prometheus.Labels{"type": "user_request_mismatch"}).Inc()
		return fmt.Errorf("Authenticated username and cert username must match")
	}

	if !session.HasScope("ca:issue") {
		caError.With(prometheus.Labels{"type": "missing_oauth_scope"}).Inc()
		return fmt.Errorf("Authorized session does not have scope ca:issue")
	}

	if certRequest.Extensions == nil {
		caError.With(prometheus.Labels{"type": "no_extensions"}).Inc()
		return fmt.Errorf("Cert request extensions are empty")
	}

	hostLine, ok := certRequest.Extensions["allowed-hosts"]
	if !ok {
		caError.With(prometheus.Labels{"type": "no_allowed_hosts"}).Inc()
		return fmt.Errorf("Cert request allowed-hosts is blank")
	}

	for _, host := range strings.Split(hostLine, ",") {
		if !user.AuthorizedForHost(host) {
			caError.With(prometheus.Labels{"type": "user_no_auth_host"}).Inc()
			return fmt.Errorf("User %s is not authorized for host %s", session.UserId, host)
		}
	}

	h.Logger.Infof("Allowing user %s to obtain SSH certificate for hosts %s", user.Username, hostLine)
	return nil
}

func (h *CAHandler) verifyRequestSignature(c *ssh.Certificate) error {
	// Copied from ssh.Certificate#bytesForSigning
	// https://cs.opensource.google/go/x/crypto/+/refs/tags/v0.11.0:ssh/certs.go;l=499-505
	c2 := *c
	c2.Signature = nil
	out := c2.Marshal()
	// Drop trailing signature length.
	return c.Verify(out[:len(out)-4], c.Signature)
}

func (h *CAHandler) HandleIssue(c echo.Context) error {
	req, err := io.ReadAll(c.Request().Body)
	if err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Unable to read request body",
		})
	}

	pubkey, _, _, _, err := ssh.ParseAuthorizedKey(req)
	if err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Error parsing certificate request",
		})
	}

	certRequest, ok := pubkey.(*ssh.Certificate)
	if !ok {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Invalid format for certificate request",
		})
	}

	if certRequest.CertType != ssh.UserCert {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "This CA only issues user certificates",
		})
	}

	if len(certRequest.ValidPrincipals) != 1 {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Invalid number of principals specified",
		})
	}

	// Kinda silly I guess but at least proves that the requestor
	// is in posession of the private key that we're signing
	if err := h.verifyRequestSignature(certRequest); err != nil {
		h.Logger.Error(err)
		return c.JSON(http.StatusUnauthorized, map[string]string{
			"error": "Invalid signature",
		})
	}

	if err := h.authorizeRequest(c, certRequest); err != nil {
		h.Logger.Error(err)
		return c.JSON(http.StatusUnauthorized, map[string]string{
			"error": "Not authorized",
		})
	}

	utcNow := time.Now().UTC()

	// Serial doesn't really matter since these are so short lived and we
	// won't be revoking them
	certToIssue := &ssh.Certificate{
		Key:             certRequest.Key,
		Serial:          uint64(utcNow.Unix()),
		CertType:        ssh.UserCert,
		KeyId:           fmt.Sprintf("%s_%d", certRequest.ValidPrincipals[0], utcNow.Unix()),
		ValidPrincipals: certRequest.ValidPrincipals,
		ValidAfter:      uint64(utcNow.Add(-5 * time.Minute).Unix()),
		ValidBefore:     uint64(utcNow.Add(h.Expiration).Unix()),
		Permissions: ssh.Permissions{
			Extensions: map[string]string{
				"permit-pty": "",
			},
		},
	}

	if err := certToIssue.SignCert(rand.Reader, h.signer); err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Error signing certificate",
		})
	}

	caSuccess.Inc()

	return c.Blob(http.StatusOK, "application/x-ssh-certificate", ssh.MarshalAuthorizedKey(certToIssue))
}