github.com/ipfans/trojan-go@v0.11.0/api/control/control.go (about) 1 package control 2 3 import ( 4 "context" 5 "encoding/json" 6 "flag" 7 "fmt" 8 "io" 9 10 "google.golang.org/grpc" 11 12 "github.com/ipfans/trojan-go/api/service" 13 "github.com/ipfans/trojan-go/common" 14 "github.com/ipfans/trojan-go/log" 15 "github.com/ipfans/trojan-go/option" 16 ) 17 18 type apiController struct { 19 address *string 20 key *string 21 hash *string 22 cert *string 23 24 cmd *string 25 password *string 26 add *bool 27 delete *bool 28 modify *bool 29 list *bool 30 uploadSpeedLimit *int 31 downloadSpeedLimit *int 32 ipLimit *int 33 ctx context.Context 34 } 35 36 func (apiController) Name() string { 37 return "api" 38 } 39 40 func (o *apiController) listUsers(apiClient service.TrojanServerServiceClient) error { 41 stream, err := apiClient.ListUsers(o.ctx, &service.ListUsersRequest{}) 42 if err != nil { 43 return err 44 } 45 defer stream.CloseSend() 46 result := []*service.ListUsersResponse{} 47 for { 48 resp, err := stream.Recv() 49 if err != nil { 50 if err == io.EOF { 51 break 52 } 53 return err 54 } 55 result = append(result, resp) 56 } 57 data, err := json.Marshal(result) 58 common.Must(err) 59 fmt.Println(string(data)) 60 return nil 61 } 62 63 func (o *apiController) getUsers(apiClient service.TrojanServerServiceClient) error { 64 stream, err := apiClient.GetUsers(o.ctx) 65 if err != nil { 66 return err 67 } 68 defer stream.CloseSend() 69 err = stream.Send(&service.GetUsersRequest{ 70 User: &service.User{ 71 Password: *o.password, 72 Hash: *o.hash, 73 }, 74 }) 75 if err != nil { 76 return err 77 } 78 resp, err := stream.Recv() 79 if err != nil { 80 return err 81 } 82 data, err := json.Marshal(resp) 83 common.Must(err) 84 fmt.Print(string(data)) 85 return nil 86 } 87 88 func (o *apiController) setUsers(apiClient service.TrojanServerServiceClient) error { 89 stream, err := apiClient.SetUsers(o.ctx) 90 if err != nil { 91 return err 92 } 93 defer stream.CloseSend() 94 95 req := &service.SetUsersRequest{ 96 Status: &service.UserStatus{ 97 User: &service.User{ 98 Password: *o.password, 99 Hash: *o.hash, 100 }, 101 IpLimit: int32(*o.ipLimit), 102 SpeedLimit: &service.Speed{ 103 UploadSpeed: uint64(*o.uploadSpeedLimit), 104 DownloadSpeed: uint64(*o.downloadSpeedLimit), 105 }, 106 }, 107 } 108 109 switch { 110 case *o.add: 111 req.Operation = service.SetUsersRequest_Add 112 case *o.modify: 113 req.Operation = service.SetUsersRequest_Modify 114 case *o.delete: 115 req.Operation = service.SetUsersRequest_Delete 116 default: 117 return common.NewError("Invalid operation") 118 } 119 120 err = stream.Send(req) 121 if err != nil { 122 return err 123 } 124 resp, err := stream.Recv() 125 if err != nil { 126 return err 127 } 128 if resp.Success { 129 fmt.Println("Done") 130 } else { 131 fmt.Println("Failed: " + resp.Info) 132 } 133 return nil 134 } 135 136 func (o *apiController) Handle() error { 137 if *o.cmd == "" { 138 return common.NewError("") 139 } 140 conn, err := grpc.Dial(*o.address, grpc.WithInsecure()) 141 if err != nil { 142 log.Error(err) 143 return nil 144 } 145 defer conn.Close() 146 apiClient := service.NewTrojanServerServiceClient(conn) 147 switch *o.cmd { 148 case "list": 149 err := o.listUsers(apiClient) 150 if err != nil { 151 log.Error(err) 152 } 153 case "get": 154 err := o.getUsers(apiClient) 155 if err != nil { 156 log.Error(err) 157 } 158 case "set": 159 err := o.setUsers(apiClient) 160 if err != nil { 161 log.Error(err) 162 } 163 default: 164 log.Error("unknown command " + *o.cmd) 165 } 166 return nil 167 } 168 169 func (o *apiController) Priority() int { 170 return 50 171 } 172 173 func init() { 174 option.RegisterHandler(&apiController{ 175 cmd: flag.String("api", "", "Connect to a Trojan-Go API service. \"-api add/get/list\""), 176 address: flag.String("api-addr", "127.0.0.1:10000", "Address of Trojan-Go API service"), 177 password: flag.String("target-password", "", "Password of the target user"), 178 hash: flag.String("target-hash", "", "Hash of the target user"), 179 add: flag.Bool("add-profile", false, "Add a new profile with API"), 180 delete: flag.Bool("delete-profile", false, "Delete an existing profile with API"), 181 modify: flag.Bool("modify-profile", false, "Modify an existing profile with API"), 182 uploadSpeedLimit: flag.Int("upload-speed-limit", 0, "Limit the upload speed with API"), 183 downloadSpeedLimit: flag.Int("download-speed-limit", 0, "Limit the download speed with API"), 184 ipLimit: flag.Int("ip-limit", 0, "Limit the number of IP with API"), 185 ctx: context.Background(), 186 }) 187 }