1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
 2  //
 3  // Licensed under the Apache License, Version 2.0 (the "License");
 4  // you may not use this file except in compliance with the License.
 5  // You may obtain a copy of the License at
 6  //
 7  //     http://www.apache.org/licenses/LICENSE-2.0
 8  //
 9  // Unless required by applicable law or agreed to in writing, software
10  // distributed under the License is distributed on an "AS IS" BASIS,
11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  // See the License for the specific language governing permissions and
13  // limitations under the License.
14  //
15  // Requirement: Any integration or derivative work must explicitly attribute
16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
17  // documentation or user interface, as detailed in the NOTICE file.
18  
19  package main
20  
21  import (
22  	"flag"
23  	"fmt"
24  	"os"
25  
26  	"github.com/Tencent/AI-Infra-Guard/common/agent"
27  	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
28  )
29  
30  func main() {
31  	var server string
32  	flag.StringVar(&server, "server", "", "server")
33  	flag.Parse()
34  	if server == "" {
35  		v := os.Getenv("AIG_SERVER")
36  		if v != "" {
37  			server = v
38  		}
39  	}
40  	if server == "" {
41  		gologger.Errorln("server is empty")
42  		return
43  	}
44  	gologger.Infoln("connect server:", server)
45  	serverUrl := fmt.Sprintf("ws://%s/api/v1/agents/ws", server)
46  	x := agent.NewAgent(agent.AgentConfig{
47  		ServerURL: serverUrl,
48  		Info: agent.AgentInfo{
49  			ID:       "test_id",
50  			HostName: "test_hostname",
51  			IP:       "127.0.0.1",
52  			Version:  "0.1",
53  			Metadata: "",
54  		},
55  	})
56  	agent2 := agent.AIInfraScanAgent{
57  		Server: server,
58  	}
59  	agent3 := agent.McpTask{Server: server}
60  	agent4 := agent.ModelRedteamReport{Server: server}
61  	agent5 := agent.AgentTask{Server: server}
62  
63  	x.RegisterTaskFunc(&agent2)
64  	x.RegisterTaskFunc(&agent3)
65  	x.RegisterTaskFunc(&agent4)
66  	x.RegisterTaskFunc(&agent5)
67  
68  	gologger.Infoln("wait task")
69  	err := x.Start()
70  	if err != nil {
71  		gologger.WithError(err).Errorln("start agent failed")
72  	}
73  	x.Stop()
74  }