159 lines
3.1 KiB
Go
159 lines
3.1 KiB
Go
package permission
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"ai-agent/internal/db"
|
|
)
|
|
|
|
type Policy string
|
|
|
|
const (
|
|
PolicyAllow Policy = "allow"
|
|
PolicyDeny Policy = "deny"
|
|
PolicyAsk Policy = "ask"
|
|
)
|
|
|
|
type Checker struct {
|
|
store *db.Store
|
|
cache map[string]Policy
|
|
mu sync.RWMutex
|
|
yolo bool
|
|
}
|
|
|
|
func NewChecker(store *db.Store, yolo bool) *Checker {
|
|
c := &Checker{
|
|
store: store,
|
|
cache: make(map[string]Policy),
|
|
yolo: yolo,
|
|
}
|
|
if store != nil {
|
|
c.loadFromDB()
|
|
}
|
|
return c
|
|
}
|
|
|
|
func (c *Checker) Check(toolName string) Policy {
|
|
if c.yolo {
|
|
return PolicyAllow
|
|
}
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
if p, ok := c.cache[toolName]; ok {
|
|
return p
|
|
}
|
|
return PolicyAsk
|
|
}
|
|
|
|
func (c *Checker) SetPolicy(toolName string, policy Policy) {
|
|
c.mu.Lock()
|
|
c.cache[toolName] = policy
|
|
c.mu.Unlock()
|
|
if c.store != nil {
|
|
c.store.UpsertToolPermission(context.Background(), db.UpsertToolPermissionParams{
|
|
ToolName: toolName,
|
|
Policy: string(policy),
|
|
})
|
|
}
|
|
}
|
|
|
|
func (c *Checker) IsYolo() bool {
|
|
return c.yolo
|
|
}
|
|
|
|
func (c *Checker) AllPolicies() map[string]Policy {
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
result := make(map[string]Policy, len(c.cache))
|
|
for k, v := range c.cache {
|
|
result[k] = v
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (c *Checker) Reset() {
|
|
c.mu.Lock()
|
|
c.cache = make(map[string]Policy)
|
|
c.mu.Unlock()
|
|
if c.store != nil {
|
|
c.store.ResetToolPermissions(context.Background())
|
|
}
|
|
}
|
|
|
|
func (c *Checker) loadFromDB() {
|
|
perms, err := c.store.ListToolPermissions(context.Background())
|
|
if err != nil {
|
|
return
|
|
}
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
for _, p := range perms {
|
|
switch Policy(p.Policy) {
|
|
case PolicyAllow, PolicyDeny, PolicyAsk:
|
|
c.cache[p.ToolName] = Policy(p.Policy)
|
|
}
|
|
}
|
|
}
|
|
|
|
type ApprovalRequest struct {
|
|
ToolName string
|
|
Args map[string]any
|
|
Response chan ApprovalResponse
|
|
}
|
|
|
|
type ApprovalResponse struct {
|
|
Allowed bool
|
|
Always bool
|
|
}
|
|
|
|
func RequestApproval(toolName string, args map[string]any, callback func(ApprovalRequest)) (bool, bool) {
|
|
if callback == nil {
|
|
return true, false
|
|
}
|
|
ch := make(chan ApprovalResponse, 1)
|
|
callback(ApprovalRequest{
|
|
ToolName: toolName,
|
|
Args: args,
|
|
Response: ch,
|
|
})
|
|
resp := <-ch
|
|
return resp.Allowed, resp.Always
|
|
}
|
|
|
|
type CheckResult int
|
|
|
|
const (
|
|
CheckAllow CheckResult = iota
|
|
CheckDeny
|
|
CheckAsk
|
|
)
|
|
|
|
func (c *Checker) ToCheckResult(toolName string) CheckResult {
|
|
if c == nil || c.yolo {
|
|
return CheckAllow
|
|
}
|
|
switch c.Check(toolName) {
|
|
case PolicyAllow:
|
|
return CheckAllow
|
|
case PolicyDeny:
|
|
return CheckDeny
|
|
default:
|
|
return CheckAsk
|
|
}
|
|
}
|
|
|
|
func NilSafe(store *db.Store, yolo bool) *Checker {
|
|
return NewChecker(store, yolo)
|
|
}
|
|
|
|
var AlwaysAllow = func(_ ApprovalRequest) {}
|
|
|
|
type ErrDenied struct {
|
|
ToolName string
|
|
}
|
|
|
|
func (e *ErrDenied) Error() string {
|
|
return "tool call denied by permission policy: " + e.ToolName
|
|
}
|