// Keep each PAM transaction on one OS thread. PAM modules can use thread-local state.
package auth

/*
#cgo LDFLAGS: -lpam
#include <security/pam_appl.h>
#include <stdlib.h>
#include <stdint.h>

// Match the signature that cgo generates for //export.
extern int latchd_pam_conv(int, struct pam_message **,
                            struct pam_response **, void *);

// Store the numeric handle in C memory. Do not pass a Go pointer to PAM.
static void *storeHandle(uintptr_t h) {
	uintptr_t *p = malloc(sizeof(uintptr_t));
	if (p) *p = h;
	return p;
}
static uintptr_t loadHandle(void *p) {
	return *(uintptr_t *)p;
}

static struct pam_conv makePAMConv(void *data) {
	struct pam_conv c;
	c.conv = (int (*)(int, const struct pam_message **,
	                  struct pam_response **, void *))latchd_pam_conv;
	c.appdata_ptr = data;
	return c;
}
*/
import "C"
import (
	"fmt"
	"os"
	ouser "os/user"
	"runtime"
	"runtime/cgo"
	"strings"
	"unsafe"
)

type AuthError struct {
	Kind int
	Msg  string
}

const (
	ErrPAMService = iota
	ErrAccountValidation
	ErrHomeDirInvalid
	ErrShellInvalid
	ErrUsernameNotFound
	ErrSessionOpen
)

func (e *AuthError) Error() string { return e.Msg }

type Credentials struct {
	Username   string
	UID        uint32
	PrimaryGID uint32
	AllGIDs    []uint32
	HomeDir    string
	Shell      string
	pamHandle  *C.pam_handle_t
	pamData    unsafe.Pointer
	pamCB      cgo.Handle
}

//export latchd_pam_conv
func latchd_pam_conv(numMsg C.int, msg **C.struct_pam_message, resp **C.struct_pam_response, data unsafe.Pointer) C.int {
	h := cgo.Handle(C.loadHandle(data))
	pw := h.Value().(string)

	r := (*C.struct_pam_response)(C.calloc(C.ulong(numMsg), C.sizeof_struct_pam_response))
	if r == nil {
		return C.PAM_BUF_ERR
	}
	*resp = r

	msgSlice := unsafe.Slice(msg, int(numMsg))
	respSlice := unsafe.Slice(r, int(numMsg))

	for i := 0; i < int(numMsg); i++ {
		switch msgSlice[i].msg_style {
		case C.PAM_PROMPT_ECHO_OFF, C.PAM_PROMPT_ECHO_ON:
			c := C.CString(pw)
			if c == nil {
				return C.PAM_BUF_ERR
			}
			respSlice[i].resp = c
			respSlice[i].resp_retcode = 0
		default:
			respSlice[i].resp = nil
			respSlice[i].resp_retcode = 0
		}
	}
	return C.PAM_SUCCESS
}

func Validate(user, pass, service string) (*Credentials, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	cSvc := C.CString(service)
	cUser := C.CString(user)
	defer C.free(unsafe.Pointer(cSvc))
	defer C.free(unsafe.Pointer(cUser))

	h := cgo.NewHandle(pass)

	cData := C.storeHandle(C.uintptr_t(h))

	conv := C.makePAMConv(cData)

	var pamh *C.pam_handle_t
	ret := C.pam_start(cSvc, cUser, &conv, &pamh)
	if ret != C.PAM_SUCCESS {
		h.Delete()
		C.free(cData)
		return nil, &AuthError{ErrPAMService, fmt.Sprintf("pam_start: %s", pamErr(pamh, ret))}
	}

	ret = C.pam_authenticate(pamh, 0)
	if ret != C.PAM_SUCCESS {
		msg := fmt.Sprintf("auth failed: %s", pamErr(pamh, ret))
		C.pam_end(pamh, ret)
		h.Delete()
		C.free(cData)
		return nil, &AuthError{ErrAccountValidation, msg}
	}

	ret = C.pam_acct_mgmt(pamh, 0)
	if ret != C.PAM_SUCCESS {
		msg := fmt.Sprintf("account: %s", pamErr(pamh, ret))
		C.pam_end(pamh, ret)
		h.Delete()
		C.free(cData)
		return nil, &AuthError{ErrAccountValidation, msg}
	}

	u, err := ouser.Lookup(user)
	if err != nil {
		C.pam_end(pamh, ret)
		h.Delete()
		C.free(cData)
		return nil, &AuthError{ErrUsernameNotFound, fmt.Sprintf("lookup %s: %v", user, err)}
	}

	uid := atou32(u.Uid)
	gid := atou32(u.Gid)
	gids := []uint32{gid}
	if ids, err := u.GroupIds(); err == nil {
		for _, s := range ids {
			g := atou32(s)
			if g != gid {
				gids = append(gids, g)
			}
		}
	}

	return &Credentials{
		Username:   user,
		UID:        uid,
		PrimaryGID: gid,
		AllGIDs:    gids,
		HomeDir:    u.HomeDir,
		Shell:      shellFromPasswd(uid),
		pamHandle:  pamh,
		pamData:    cData,
		pamCB:      h,
	}, nil
}

func (c *Credentials) OpenSession() error {
	ret := C.pam_open_session(c.pamHandle, 0)
	if ret != C.PAM_SUCCESS {
		return &AuthError{ErrSessionOpen, fmt.Sprintf("session: %s", pamErr(c.pamHandle, ret))}
	}
	return nil
}

func (c *Credentials) CloseSession() {
	if c.pamHandle != nil {
		C.pam_close_session(c.pamHandle, 0)
		C.pam_end(c.pamHandle, 0)
		c.pamHandle = nil
	}
	if c.pamData != nil {
		C.free(c.pamData)
		c.pamData = nil
	}
	c.pamCB.Delete()
}

func pamErr(h *C.pam_handle_t, e C.int) string {
	return C.GoString(C.pam_strerror(h, e))
}

func atou32(s string) uint32 {
	var n uint32
	for _, c := range s {
		if c < '0' || c > '9' {
			break
		}
		n = n*10 + uint32(c-'0')
	}
	return n
}

func shellFromPasswd(uid uint32) string {
	data, err := os.ReadFile("/etc/passwd")
	if err != nil {
		return "/bin/sh"
	}
	uidStr := fmt.Sprint(uid)
	for _, line := range strings.Split(string(data), "\n") {
		parts := strings.Split(line, ":")
		if len(parts) >= 7 && parts[2] == uidStr && parts[6] != "" {
			sh := parts[6]
			if sh != "/sbin/nologin" && sh != "/usr/sbin/nologin" && sh != "/bin/false" {
				return sh
			}
		}
	}
	return "/bin/sh"
}