first commit
Agent Release / build (push) Has been cancelled
Server Deploy / deploy (push) Has been cancelled

This commit is contained in:
domrichardson
2026-06-15 13:58:45 +01:00
commit c9868b2108
55 changed files with 11076 additions and 0 deletions
+293
View File
@@ -0,0 +1,293 @@
package api
import (
"fmt"
"net/http"
"os"
"github.com/gin-gonic/gin"
"github.com/mrhid6/keymanager/server/internal/models"
"github.com/mrhid6/keymanager/server/internal/services"
)
func RegisterRoutes(r *gin.Engine) {
r.GET("/install", handleInstallScript)
api := r.Group("/api")
{
api.GET("/servers", listServers)
api.POST("/servers", createServer)
api.GET("/servers/new", newServer)
api.POST("/servers/new", newServer)
api.GET("/servers/:id", getServer)
api.DELETE("/servers/:id", deleteServer)
api.POST("/servers/:id/generate-key", generateKey)
api.GET("/keys", listKeys)
api.POST("/keys", createKey)
api.GET("/keys/:id", getKey)
api.POST("/keys/:id/assign", assignKey)
api.DELETE("/keys/:id/assign/:serverId", revokeAssignment)
}
}
func listServers(c *gin.Context) {
servers, err := services.ListServers()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, servers)
}
func createServer(c *gin.Context) {
s, token, err := services.CreateServer()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusCreated, gin.H{
"server": s,
"token": token,
"server_id": s.ServerID,
})
}
func newServer(c *gin.Context) {
s, token, err := services.CreateServer()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
giteaHost := os.Getenv("GITEA_HOST")
if giteaHost == "" {
giteaHost = "gitea.example.com"
}
host := os.Getenv("PUBLIC_HOST")
if host == "" {
host = "keymanager.example.com"
}
installCmd := fmt.Sprintf(
`curl -fsSL "https://%s/install?server_id=%s&token=%s" | bash`,
host, s.ServerID, token,
)
c.JSON(http.StatusOK, gin.H{
"server_id": s.ServerID,
"pre_reg_token": token,
"install_command": installCmd,
})
}
func getServer(c *gin.Context) {
id := c.Param("id")
s, err := services.GetServer(id)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
return
}
assignments, _ := services.GetAssignmentsWithKeysForServer(id)
// Build response matching ServerWithKeys shape expected by frontend
type serverResponse struct {
*models.Server
Keys interface{} `json:"keys"`
}
c.JSON(http.StatusOK, serverResponse{
Server: s,
Keys: assignments,
})
}
func deleteServer(c *gin.Context) {
id := c.Param("id")
if err := services.DeleteServer(id); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"deleted": true})
}
func generateKey(c *gin.Context) {
// The agent triggers key generation itself; this endpoint signals
// the intent by returning the server so the caller knows to wait
// for the agent to upload via gRPC UploadGeneratedKey.
id := c.Param("id")
s, err := services.GetServer(id)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
return
}
c.JSON(http.StatusOK, gin.H{
"message": "agent will generate and upload key on next poll",
"server_id": s.ServerID,
})
}
func listKeys(c *gin.Context) {
keys, err := services.ListKeys()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, keys)
}
func createKey(c *gin.Context) {
var body struct {
Label string `json:"label" binding:"required"`
PublicKey string `json:"public_key" binding:"required"`
}
if err := c.ShouldBindJSON(&body); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
key, err := services.CreateKey(body.Label, body.PublicKey, "uploaded", "")
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusCreated, key)
}
func getKey(c *gin.Context) {
id := c.Param("id")
key, err := services.GetKey(id)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "key not found"})
return
}
assignments, _ := services.GetAssignmentsWithServers(id)
c.JSON(http.StatusOK, gin.H{
"key": key,
"assignments": assignments,
})
}
func assignKey(c *gin.Context) {
keyID := c.Param("id")
var body struct {
ServerID string `json:"server_id" binding:"required"`
}
if err := c.ShouldBindJSON(&body); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
a, err := services.AssignKey(keyID, body.ServerID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusCreated, a)
}
func revokeAssignment(c *gin.Context) {
keyID := c.Param("id")
serverID := c.Param("serverId")
if err := services.RevokeAssignment(keyID, serverID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"revoked": true})
}
func handleInstallScript(c *gin.Context) {
serverID := c.Query("server_id")
token := c.Query("token")
giteaHost := os.Getenv("GITEA_HOST")
if giteaHost == "" {
giteaHost = "gitea.example.com"
}
publicHost := os.Getenv("PUBLIC_HOST")
if publicHost == "" {
publicHost = "keymanager.example.com"
}
script := fmt.Sprintf(`#!/usr/bin/env bash
set -euo pipefail
SERVER_ID="%s"
TOKEN="%s"
GITEA_HOST="%s"
KM_HOST="%s"
ARCH=$(uname -m)
case "$ARCH" in
x86_64) ARCH="amd64" ;;
aarch64) ARCH="arm64" ;;
*) echo "Unsupported architecture: $ARCH" >&2; exit 1 ;;
esac
# Get latest agent release tag
LATEST=$(curl -fsSL "https://${GITEA_HOST}/api/v1/repos/mrhid6/keymanager/releases?limit=10" \
| grep -o '"tag_name":"agent/v[^"]*"' | head -1 | sed 's/"tag_name":"//;s/"//')
if [ -z "$LATEST" ]; then
echo "Could not determine latest agent version" >&2
exit 1
fi
VERSION="${LATEST#agent/}"
BINARY_URL="https://${GITEA_HOST}/mrhid6/keymanager/releases/download/${LATEST}/keymanager-agent-linux-${ARCH}"
CHECKSUM_URL="https://${GITEA_HOST}/mrhid6/keymanager/releases/download/${LATEST}/checksums.txt"
echo "Installing keymanager-agent ${VERSION} (${ARCH})..."
curl -fsSL -o /tmp/keymanager-agent "${BINARY_URL}"
curl -fsSL -o /tmp/checksums.txt "${CHECKSUM_URL}"
cd /tmp
EXPECTED=$(grep "keymanager-agent-linux-${ARCH}" checksums.txt | awk '{print $1}')
ACTUAL=$(sha256sum keymanager-agent | awk '{print $1}')
if [ "$EXPECTED" != "$ACTUAL" ]; then
echo "Checksum mismatch!" >&2
exit 1
fi
install -m 0755 /tmp/keymanager-agent /usr/local/bin/keymanager-agent
mkdir -p /etc/keymanager
chmod 0700 /etc/keymanager
cat > /etc/keymanager/config.yaml <<EOF
server_url: "${KM_HOST}:9090"
server_id: "${SERVER_ID}"
pre_reg_token: "${TOKEN}"
agent_token: ""
poll_interval: 30s
tls: true
EOF
chmod 0600 /etc/keymanager/config.yaml
cat > /etc/systemd/system/keymanager-agent.service <<EOF
[Unit]
Description=KeyManager Agent
After=network.target
[Service]
ExecStart=/usr/local/bin/keymanager-agent
Restart=always
RestartSec=10
User=root
[Install]
WantedBy=multi-user.target
EOF
systemctl daemon-reload
systemctl enable --now keymanager-agent
echo "keymanager-agent installed and started."
`, serverID, token, giteaHost, publicHost)
c.Header("Content-Type", "text/x-shellscript")
c.String(http.StatusOK, script)
}
+34
View File
@@ -0,0 +1,34 @@
package db
import (
"context"
"time"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
)
var Client *mongo.Client
var Database *mongo.Database
func Connect(uri, dbName string) error {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
client, err := mongo.Connect(options.Client().ApplyURI(uri))
if err != nil {
return err
}
if err := client.Ping(ctx, nil); err != nil {
return err
}
Client = client
Database = client.Database(dbName)
return nil
}
func Col(name string) *mongo.Collection {
return Database.Collection(name)
}
+20
View File
@@ -0,0 +1,20 @@
package grpcserver
import (
"encoding/json"
)
// JSONCodec is a gRPC codec that uses JSON encoding.
type JSONCodec struct{}
func (JSONCodec) Marshal(v interface{}) ([]byte, error) {
return json.Marshal(v)
}
func (JSONCodec) Unmarshal(data []byte, v interface{}) error {
return json.Unmarshal(data, v)
}
func (JSONCodec) Name() string {
return "proto" // override default proto codec name so gRPC uses it
}
+171
View File
@@ -0,0 +1,171 @@
// Hand-written gRPC bindings for keymanager.proto using JSON codec.
// To use: register the JSON codec before creating gRPC servers/clients.
package pb
import (
"context"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
// Message types
type RegisterRequest struct {
ServerId string `json:"server_id"`
PreRegToken string `json:"pre_reg_token"`
Hostname string `json:"hostname"`
IpAddress string `json:"ip_address"`
OsInfo string `json:"os_info"`
}
type RegisterResponse struct {
AgentToken string `json:"agent_token"`
}
type SyncRequest struct {
ServerId string `json:"server_id"`
AgentToken string `json:"agent_token"`
}
type SyncResponse struct {
PublicKeys []string `json:"public_keys"`
}
type UploadKeyRequest struct {
ServerId string `json:"server_id"`
AgentToken string `json:"agent_token"`
PublicKey string `json:"public_key"`
Label string `json:"label"`
}
type UploadKeyResponse struct {
KeyId string `json:"key_id"`
}
// Server interface
type KeyManagerServer interface {
Register(context.Context, *RegisterRequest) (*RegisterResponse, error)
SyncKeys(context.Context, *SyncRequest) (*SyncResponse, error)
UploadGeneratedKey(context.Context, *UploadKeyRequest) (*UploadKeyResponse, error)
}
type UnimplementedKeyManagerServer struct{}
func (UnimplementedKeyManagerServer) Register(context.Context, *RegisterRequest) (*RegisterResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Register not implemented")
}
func (UnimplementedKeyManagerServer) SyncKeys(context.Context, *SyncRequest) (*SyncResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method SyncKeys not implemented")
}
func (UnimplementedKeyManagerServer) UploadGeneratedKey(context.Context, *UploadKeyRequest) (*UploadKeyResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method UploadGeneratedKey not implemented")
}
// Client interface
type KeyManagerClient interface {
Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error)
}
type keyManagerClient struct {
cc grpc.ClientConnInterface
}
func NewKeyManagerClient(cc grpc.ClientConnInterface) KeyManagerClient {
return &keyManagerClient{cc}
}
func (c *keyManagerClient) Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error) {
out := new(RegisterResponse)
if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/Register", in, out, opts...); err != nil {
return nil, err
}
return out, nil
}
func (c *keyManagerClient) SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error) {
out := new(SyncResponse)
if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/SyncKeys", in, out, opts...); err != nil {
return nil, err
}
return out, nil
}
func (c *keyManagerClient) UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error) {
out := new(UploadKeyResponse)
if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/UploadGeneratedKey", in, out, opts...); err != nil {
return nil, err
}
return out, nil
}
// Server registration
func RegisterKeyManagerServer(s grpc.ServiceRegistrar, srv KeyManagerServer) {
s.RegisterService(&KeyManager_ServiceDesc, srv)
}
var KeyManager_ServiceDesc = grpc.ServiceDesc{
ServiceName: "keymanager.v1.KeyManager",
HandlerType: (*KeyManagerServer)(nil),
Methods: []grpc.MethodDesc{
{MethodName: "Register", Handler: _KeyManager_Register_Handler},
{MethodName: "SyncKeys", Handler: _KeyManager_SyncKeys_Handler},
{MethodName: "UploadGeneratedKey", Handler: _KeyManager_UploadGeneratedKey_Handler},
},
Streams: []grpc.StreamDesc{},
Metadata: "keymanager/v1/keymanager.proto",
}
func _KeyManager_Register_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(RegisterRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(KeyManagerServer).Register(ctx, in)
}
info := &grpc.UnaryServerInfo{Server: srv, FullMethod: "/keymanager.v1.KeyManager/Register"}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(KeyManagerServer).Register(ctx, req.(*RegisterRequest))
}
return interceptor(ctx, in, info, handler)
}
func _KeyManager_SyncKeys_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(SyncRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(KeyManagerServer).SyncKeys(ctx, in)
}
info := &grpc.UnaryServerInfo{Server: srv, FullMethod: "/keymanager.v1.KeyManager/SyncKeys"}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(KeyManagerServer).SyncKeys(ctx, req.(*SyncRequest))
}
return interceptor(ctx, in, info, handler)
}
func _KeyManager_UploadGeneratedKey_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(UploadKeyRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(KeyManagerServer).UploadGeneratedKey(ctx, in)
}
info := &grpc.UnaryServerInfo{Server: srv, FullMethod: "/keymanager.v1.KeyManager/UploadGeneratedKey"}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(KeyManagerServer).UploadGeneratedKey(ctx, req.(*UploadKeyRequest))
}
return interceptor(ctx, in, info, handler)
}
+81
View File
@@ -0,0 +1,81 @@
package grpcserver
import (
"context"
"fmt"
"log"
"net"
"github.com/mrhid6/keymanager/server/internal/grpc/pb"
"github.com/mrhid6/keymanager/server/internal/services"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/encoding"
"google.golang.org/grpc/status"
)
func init() {
encoding.RegisterCodec(JSONCodec{})
}
type keyManagerServer struct {
pb.UnimplementedKeyManagerServer
}
func (s *keyManagerServer) Register(ctx context.Context, req *pb.RegisterRequest) (*pb.RegisterResponse, error) {
agentToken, err := services.RegisterServer(req.ServerId, req.PreRegToken, req.Hostname, req.IpAddress, req.OsInfo)
if err != nil {
return nil, status.Errorf(codes.InvalidArgument, "registration failed: %v", err)
}
return &pb.RegisterResponse{AgentToken: agentToken}, nil
}
func (s *keyManagerServer) SyncKeys(ctx context.Context, req *pb.SyncRequest) (*pb.SyncResponse, error) {
srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
if err != nil {
return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
}
if err := services.UpdateServerLastSeen(srv.ServerID); err != nil {
log.Printf("failed to update last seen for %s: %v", srv.ServerID, err)
}
keys, err := services.BuildAuthorizedKeys(req.ServerId)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to build authorized keys: %v", err)
}
return &pb.SyncResponse{PublicKeys: keys}, nil
}
func (s *keyManagerServer) UploadGeneratedKey(ctx context.Context, req *pb.UploadKeyRequest) (*pb.UploadKeyResponse, error) {
srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
if err != nil {
return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
}
key, err := services.CreateKey(req.Label, req.PublicKey, "generated", srv.ServerID)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to store key: %v", err)
}
// Auto-assign to the generating server
if _, err := services.AssignKey(key.KeyID, srv.ServerID); err != nil {
log.Printf("failed to auto-assign generated key: %v", err)
}
return &pb.UploadKeyResponse{KeyId: key.KeyID}, nil
}
func StartGRPC(port int) error {
lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
if err != nil {
return fmt.Errorf("failed to listen: %w", err)
}
s := grpc.NewServer()
pb.RegisterKeyManagerServer(s, &keyManagerServer{})
log.Printf("gRPC server listening on :%d", port)
return s.Serve(lis)
}
+15
View File
@@ -0,0 +1,15 @@
package models
import (
"time"
"go.mongodb.org/mongo-driver/v2/bson"
)
type Assignment struct {
ID bson.ObjectID `bson:"_id,omitempty" json:"_id,omitempty"`
KeyID string `bson:"key_id" json:"key_id"`
ServerID string `bson:"server_id" json:"server_id"`
AssignedAt time.Time `bson:"assigned_at" json:"assigned_at"`
RevokedAt *time.Time `bson:"revoked_at,omitempty" json:"revoked_at,omitempty"`
}
+18
View File
@@ -0,0 +1,18 @@
package models
import (
"time"
"go.mongodb.org/mongo-driver/v2/bson"
)
type Key struct {
ID bson.ObjectID `bson:"_id,omitempty" json:"_id,omitempty"`
KeyID string `bson:"key_id" json:"key_id"`
Label string `bson:"label" json:"label"`
PublicKey string `bson:"public_key" json:"public_key"`
Fingerprint string `bson:"fingerprint" json:"fingerprint"`
Source string `bson:"source" json:"source"` // uploaded | generated
GeneratedByServerID string `bson:"generated_by_server_id,omitempty" json:"generated_by_server_id,omitempty"`
CreatedAt time.Time `bson:"created_at" json:"created_at"`
}
+21
View File
@@ -0,0 +1,21 @@
package models
import (
"time"
"go.mongodb.org/mongo-driver/v2/bson"
)
type Server struct {
ID bson.ObjectID `bson:"_id,omitempty" json:"_id,omitempty"`
ServerID string `bson:"server_id" json:"server_id"`
Hostname string `bson:"hostname" json:"hostname"`
IPAddress string `bson:"ip_address" json:"ip_address"`
OSInfo string `bson:"os_info" json:"os_info"`
PreRegToken string `bson:"pre_reg_token,omitempty" json:"pre_reg_token,omitempty"`
PreRegExpires *time.Time `bson:"pre_reg_expires,omitempty" json:"pre_reg_expires,omitempty"`
AgentTokenHash string `bson:"agent_token_hash,omitempty" json:"-"`
Status string `bson:"status" json:"status"`
LastSeen *time.Time `bson:"last_seen,omitempty" json:"last_seen,omitempty"`
CreatedAt time.Time `bson:"created_at" json:"created_at"`
}
+209
View File
@@ -0,0 +1,209 @@
package services
import (
"context"
"crypto/md5"
"encoding/base64"
"fmt"
"strings"
"time"
"github.com/google/uuid"
"github.com/mrhid6/keymanager/server/internal/db"
"github.com/mrhid6/keymanager/server/internal/models"
"go.mongodb.org/mongo-driver/v2/bson"
)
func computeFingerprint(pubKey string) string {
parts := strings.Fields(pubKey)
if len(parts) < 2 {
return ""
}
raw, err := base64.StdEncoding.DecodeString(parts[1])
if err != nil {
return ""
}
sum := md5.Sum(raw)
var pairs []string
for _, b := range sum {
pairs = append(pairs, fmt.Sprintf("%02x", b))
}
return "MD5:" + strings.Join(pairs, ":")
}
func CreateKey(label, publicKey, source, generatedByServerID string) (*models.Key, error) {
key := &models.Key{
KeyID: uuid.NewString(),
Label: label,
PublicKey: publicKey,
Fingerprint: computeFingerprint(publicKey),
Source: source,
GeneratedByServerID: generatedByServerID,
CreatedAt: time.Now(),
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err := db.Col("keys").InsertOne(ctx, key)
if err != nil {
return nil, err
}
return key, nil
}
func GetKey(keyID string) (*models.Key, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var key models.Key
err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID}).Decode(&key)
if err != nil {
return nil, err
}
return &key, nil
}
func ListKeys() ([]models.Key, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cursor, err := db.Col("keys").Find(ctx, bson.M{})
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var keys []models.Key
if err := cursor.All(ctx, &keys); err != nil {
return nil, err
}
return keys, nil
}
func DeleteKey(keyID string) error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err := db.Col("keys").DeleteOne(ctx, bson.M{"key_id": keyID})
return err
}
func AssignKey(keyID, serverID string) (*models.Assignment, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Check if already assigned and active
var existing models.Assignment
err := db.Col("assignments").FindOne(ctx, bson.M{
"key_id": keyID,
"server_id": serverID,
"revoked_at": nil,
}).Decode(&existing)
if err == nil {
return &existing, nil
}
a := &models.Assignment{
KeyID: keyID,
ServerID: serverID,
AssignedAt: time.Now(),
}
_, err = db.Col("assignments").InsertOne(ctx, a)
if err != nil {
return nil, err
}
return a, nil
}
func RevokeAssignment(keyID, serverID string) error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
now := time.Now()
_, err := db.Col("assignments").UpdateOne(ctx,
bson.M{"key_id": keyID, "server_id": serverID, "revoked_at": nil},
bson.M{"$set": bson.M{"revoked_at": now}},
)
return err
}
func GetAssignmentsForKey(keyID string) ([]models.Assignment, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cursor, err := db.Col("assignments").Find(ctx, bson.M{"key_id": keyID, "revoked_at": nil})
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var assignments []models.Assignment
if err := cursor.All(ctx, &assignments); err != nil {
return nil, err
}
return assignments, nil
}
type AssignmentWithServer struct {
models.Assignment `bson:",inline"`
Server *models.Server `json:"server,omitempty"`
}
func GetAssignmentsWithServers(keyID string) ([]AssignmentWithServer, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cursor, err := db.Col("assignments").Find(ctx, bson.M{"key_id": keyID})
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var assignments []models.Assignment
if err := cursor.All(ctx, &assignments); err != nil {
return nil, err
}
result := make([]AssignmentWithServer, 0, len(assignments))
for _, a := range assignments {
item := AssignmentWithServer{Assignment: a}
var srv models.Server
if err := db.Col("servers").FindOne(ctx, bson.M{"server_id": a.ServerID}).Decode(&srv); err == nil {
item.Server = &srv
}
result = append(result, item)
}
return result, nil
}
type AssignmentWithKey struct {
models.Assignment `bson:",inline"`
Key *models.Key `json:"key,omitempty"`
}
func GetAssignmentsWithKeysForServer(serverID string) ([]AssignmentWithKey, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cursor, err := db.Col("assignments").Find(ctx, bson.M{"server_id": serverID})
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var assignments []models.Assignment
if err := cursor.All(ctx, &assignments); err != nil {
return nil, err
}
result := make([]AssignmentWithKey, 0, len(assignments))
for _, a := range assignments {
item := AssignmentWithKey{Assignment: a}
var key models.Key
if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": a.KeyID}).Decode(&key); err == nil {
item.Key = &key
}
result = append(result, item)
}
return result, nil
}
+193
View File
@@ -0,0 +1,193 @@
package services
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/hex"
"fmt"
"time"
"github.com/google/uuid"
"github.com/mrhid6/keymanager/server/internal/db"
"github.com/mrhid6/keymanager/server/internal/models"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo/options"
)
func generateToken(n int) (string, error) {
b := make([]byte, n)
if _, err := rand.Read(b); err != nil {
return "", err
}
return hex.EncodeToString(b), nil
}
func HashToken(token string) string {
sum := sha256.Sum256([]byte(token))
return hex.EncodeToString(sum[:])
}
func CreateServer() (*models.Server, string, error) {
token, err := generateToken(32)
if err != nil {
return nil, "", err
}
expires := time.Now().Add(time.Hour)
s := &models.Server{
ServerID: uuid.NewString(),
PreRegToken: token,
PreRegExpires: &expires,
Status: "pending",
CreatedAt: time.Now(),
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err = db.Col("servers").InsertOne(ctx, s)
if err != nil {
return nil, "", err
}
return s, token, nil
}
func GetServer(serverID string) (*models.Server, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var s models.Server
err := db.Col("servers").FindOne(ctx, bson.M{"server_id": serverID}).Decode(&s)
if err != nil {
return nil, err
}
return &s, nil
}
func GetServerByPreRegToken(token string) (*models.Server, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var s models.Server
err := db.Col("servers").FindOne(ctx, bson.M{
"pre_reg_token": token,
"pre_reg_expires": bson.M{"$gt": time.Now()},
}).Decode(&s)
if err != nil {
return nil, err
}
return &s, nil
}
func RegisterServer(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var s models.Server
err := db.Col("servers").FindOne(ctx, bson.M{
"server_id": serverID,
"pre_reg_token": preRegToken,
"pre_reg_expires": bson.M{"$gt": time.Now()},
}).Decode(&s)
if err != nil {
return "", fmt.Errorf("invalid or expired pre-registration token")
}
agentToken, err := generateToken(32)
if err != nil {
return "", err
}
tokenHash := HashToken(agentToken)
now := time.Now()
_, err = db.Col("servers").UpdateOne(ctx,
bson.M{"server_id": serverID},
bson.M{"$set": bson.M{
"hostname": hostname,
"ip_address": ipAddress,
"os_info": osInfo,
"agent_token_hash": tokenHash,
"status": "active",
"last_seen": now,
"pre_reg_token": "",
"pre_reg_expires": nil,
}},
)
if err != nil {
return "", err
}
return agentToken, nil
}
func ValidateAgentToken(serverID, agentToken string) (*models.Server, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
tokenHash := HashToken(agentToken)
var s models.Server
err := db.Col("servers").FindOne(ctx, bson.M{
"server_id": serverID,
"agent_token_hash": tokenHash,
}).Decode(&s)
if err != nil {
return nil, fmt.Errorf("invalid agent token")
}
return &s, nil
}
func UpdateServerLastSeen(serverID string) error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
now := time.Now()
_, err := db.Col("servers").UpdateOne(ctx,
bson.M{"server_id": serverID},
bson.M{"$set": bson.M{"last_seen": now, "status": "active"}},
)
return err
}
func ListServers() ([]models.Server, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}})
cursor, err := db.Col("servers").Find(ctx, bson.M{}, opts)
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var servers []models.Server
if err := cursor.All(ctx, &servers); err != nil {
return nil, err
}
return servers, nil
}
func DeleteServer(serverID string) error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err := db.Col("servers").DeleteOne(ctx, bson.M{"server_id": serverID})
if err != nil {
return err
}
// Also remove assignments
_, err = db.Col("assignments").DeleteMany(ctx, bson.M{"server_id": serverID})
return err
}
func MarkOfflineServers(threshold time.Duration) error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cutoff := time.Now().Add(-threshold)
_, err := db.Col("servers").UpdateMany(ctx,
bson.M{
"status": "active",
"last_seen": bson.M{"$lt": cutoff},
},
bson.M{"$set": bson.M{"status": "offline"}},
)
return err
}
+40
View File
@@ -0,0 +1,40 @@
package services
import (
"context"
"time"
"github.com/mrhid6/keymanager/server/internal/db"
"github.com/mrhid6/keymanager/server/internal/models"
"go.mongodb.org/mongo-driver/v2/bson"
)
func BuildAuthorizedKeys(serverID string) ([]string, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cursor, err := db.Col("assignments").Find(ctx, bson.M{
"server_id": serverID,
"revoked_at": nil,
})
if err != nil {
return nil, err
}
defer cursor.Close(ctx)
var assignments []models.Assignment
if err := cursor.All(ctx, &assignments); err != nil {
return nil, err
}
var lines []string
for _, a := range assignments {
var key models.Key
err := db.Col("keys").FindOne(ctx, bson.M{"key_id": a.KeyID}).Decode(&key)
if err != nil {
continue
}
lines = append(lines, key.PublicKey)
}
return lines, nil
}