termcloud/internal/db/policy.go

package db

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

type BucketPolicy struct {
	Version   string            `json:"Version"`
	Statement []PolicyStatement `json:"Statement"`
}

type PolicyStatement struct {
	Sid       string         `json:"Sid,omitempty"`
	Effect    string         `json:"Effect"`
	Principal map[string]any `json:"Principal,omitempty"`
	Action    []string       `json:"Action"`
	Resource  []string       `json:"Resource"`
	Condition map[string]any `json:"Condition,omitempty"`
}

type BucketPolicyRecord struct {
	ID        int64     `json:"id"`
	BucketID  int64     `json:"bucketId"`
	Policy    string    `json:"policy"`
	UpdatedAt time.Time `json:"updatedAt"`
}

func ValidateBucketPolicy(policyJSON string) (*BucketPolicy, error) {
	var policy BucketPolicy
	if err := json.Unmarshal([]byte(policyJSON), &policy); err != nil {
		return nil, fmt.Errorf("invalid JSON format: %w", err)
	}

	if policy.Version == "" {
		policy.Version = "2012-10-17"
	}

	if len(policy.Statement) == 0 {
		return nil, fmt.Errorf("policy must contain at least one statement")
	}

	for i, stmt := range policy.Statement {
		if err := validateStatement(stmt); err != nil {
			return nil, fmt.Errorf("statement %d: %w", i, err)
		}
	}

	return &policy, nil
}

func validateStatement(stmt PolicyStatement) error {
	if stmt.Effect != "Allow" && stmt.Effect != "Deny" {
		return fmt.Errorf("effect must be 'Allow' or 'Deny'")
	}

	if len(stmt.Action) == 0 {
		return fmt.Errorf("statement must specify at least one action")
	}

	validActions := map[string]bool{
		"termcloud:GetObject":    true,
		"termcloud:PutObject":    true,
		"termcloud:DeleteObject": true,
		"termcloud:ListObjects":  true,
		"termcloud:GetBucket":    true,
		"termcloud:DeleteBucket": true,
		"*":                      true,
	}

	for _, action := range stmt.Action {
		if !validActions[action] && !strings.HasSuffix(action, "*") {
			return fmt.Errorf("invalid action: %s", action)
		}
	}

	if len(stmt.Resource) == 0 {
		return fmt.Errorf("statement must specify at least one resource")
	}

	return nil
}

func (s *BucketService) SetBucketPolicy(ctx context.Context, bucketID int64, policyJSON string) error {
	if _, err := ValidateBucketPolicy(policyJSON); err != nil {
		return fmt.Errorf("invalid policy: %w", err)
	}

	_, err := s.pool.Exec(ctx, `
		INSERT INTO bucket_policies (bucket_id, policy) 
		VALUES ($1, $2)
		ON CONFLICT (bucket_id) 
		DO UPDATE SET policy = $2, updated_at = NOW()`,
		bucketID, policyJSON)
	if err != nil {
		return fmt.Errorf("failed to set bucket policy: %w", err)
	}

	return nil
}

func (s *BucketService) GetBucketPolicy(ctx context.Context, bucketID int64) (*BucketPolicyRecord, error) {
	var record BucketPolicyRecord
	err := s.pool.QueryRow(ctx, `
		SELECT id, bucket_id, policy, updated_at 
		FROM bucket_policies WHERE bucket_id = $1`,
		bucketID).Scan(&record.ID, &record.BucketID, &record.Policy, &record.UpdatedAt)
	if err != nil {
		return nil, fmt.Errorf("failed to get bucket policy: %w", err)
	}

	return &record, nil
}

func (s *BucketService) DeleteBucketPolicy(ctx context.Context, bucketID int64) error {
	_, err := s.pool.Exec(ctx, "DELETE FROM bucket_policies WHERE bucket_id = $1", bucketID)
	if err != nil {
		return fmt.Errorf("failed to delete bucket policy: %w", err)
	}

	return nil
}

func (s *BucketService) EvaluatePolicy(ctx context.Context, bucketID int64, action, resource, principal string) (bool, error) {
	policyRecord, err := s.GetBucketPolicy(ctx, bucketID)
	if err != nil {
		if strings.Contains(err.Error(), "no rows") {
			return true, nil
		}
		return false, err
	}

	policy, err := ValidateBucketPolicy(policyRecord.Policy)
	if err != nil {
		return false, fmt.Errorf("invalid stored policy: %w", err)
	}

	for _, stmt := range policy.Statement {
		if matchesStatement(stmt, action, resource, principal) {
			return stmt.Effect == "Allow", nil
		}
	}

	return true, nil
}

func matchesStatement(stmt PolicyStatement, action, resource, principal string) bool {
	if !matchesAction(stmt.Action, action) {
		return false
	}

	if !matchesResource(stmt.Resource, resource) {
		return false
	}

	if stmt.Principal != nil && !matchesPrincipal(stmt.Principal, principal) {
		return false
	}

	return true
}

func matchesAction(allowedActions []string, action string) bool {
	for _, allowed := range allowedActions {
		if allowed == "*" || allowed == action {
			return true
		}
		if strings.HasSuffix(allowed, "*") {
			prefix := strings.TrimSuffix(allowed, "*")
			if strings.HasPrefix(action, prefix) {
				return true
			}
		}
	}
	return false
}

func matchesResource(allowedResources []string, resource string) bool {
	for _, allowed := range allowedResources {
		if allowed == "*" || allowed == resource {
			return true
		}
		if strings.HasSuffix(allowed, "*") {
			prefix := strings.TrimSuffix(allowed, "*")
			if strings.HasPrefix(resource, prefix) {
				return true
			}
		}
	}
	return false
}

func matchesPrincipal(allowedPrincipal map[string]any, principal string) bool {
	if allowedPrincipal["*"] != nil {
		return true
	}

	if users, ok := allowedPrincipal["User"].([]any); ok {
		for _, user := range users {
			if userStr, ok := user.(string); ok && userStr == principal {
				return true
			}
		}
	}

	if userStr, ok := allowedPrincipal["User"].(string); ok && userStr == principal {
		return true
	}

	return false
}