diff --git a/script/script.go b/script/script.go index 5513c9e516096fc6cc7f77ea6537c5d4f4aa7ea0..72e8076e15fb976c115f7a364a2d13b259a21c64 100644 --- a/script/script.go +++ b/script/script.go @@ -7,6 +7,7 @@ import ( "sync" "github.com/d5/tengo/script" + "github.com/d5/tengo/stdlib" ippb "git.autistici.org/ai3/tools/iprep/proto" ) @@ -24,7 +25,7 @@ func LoadScript(path string) (*Script, error) { } var globalVars = []string{ - "ip", "score", "counts", "interval", + "ip", "score", "counts", "interval", "ext", } func NewScript(src []byte) (*Script, error) { @@ -36,6 +37,10 @@ func NewScript(src []byte) (*Script, error) { return nil, err } } + + // Add available imports. + s.SetImports(stdlib.GetModuleMap("math", "times")) + c, err := s.Compile() if err != nil { return nil, err @@ -53,11 +58,12 @@ func buildIPMap(ip string, m ippb.Map) map[string]interface{} { return out } -func (script *Script) RunIP(ctx context.Context, ip string, m ippb.Map, intervalSecs float64) (float64, error) { +func (script *Script) RunIP(ctx context.Context, ip string, m ippb.Map, intervalSecs float64, ext map[string]interface{}) (float64, error) { c := script.compiled.Clone() c.Set("ip", ip) c.Set("score", 0.0) c.Set("interval", intervalSecs) + c.Set("ext", ext) if err := c.Set("counts", buildIPMap(ip, m)); err != nil { return 0, err } diff --git a/script/script_test.go b/script/script_test.go index a60fea1106c14f5e98294a2df0e2f06cd9b30276..0ddc89e243a59a7b6b63db33cc21cba104bcc428 100644 --- a/script/script_test.go +++ b/script/script_test.go @@ -22,7 +22,7 @@ score = counts["test"] / 2 + counts["test2"] m.Incr("test2", "1.2.3.4", 2) m.Incr("test2", "2.3.4.5", 3) - score, err := s.RunIP(context.Background(), "1.2.3.4", m, 3600) + score, err := s.RunIP(context.Background(), "1.2.3.4", m, 3600, nil) if err != nil { t.Fatalf("runScript: %v", err) } diff --git a/server/server.go b/server/server.go index 86f0ef93b184f1ae56b4e475163b4dfe1b455452..cf2a91985f9a01ae3acb14c7858bf818facac81f 100644 --- a/server/server.go +++ b/server/server.go @@ -88,7 +88,7 @@ func (s *Server) GetScore(ctx context.Context, req *ippb.GetScoreRequest) (*ippb return nil, status.Errorf(codes.Unavailable, "%v", err) } - score, err := s.Script().RunIP(ctx, req.Ip, m, s.horizon.Seconds()) + score, err := s.Script().RunIP(ctx, req.Ip, m, s.horizon.Seconds(), nil) if err != nil { return nil, status.Errorf(codes.Internal, "%v", err) }