add api controller, build tags

api/control/control.go

@@ -0,0 +1,181 @@
+package control
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"github.com/p4gefau1t/trojan-go/api/service"
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/log"
+	"github.com/p4gefau1t/trojan-go/option"
+	"google.golang.org/grpc"
+	"io"
+)
+
+type apiController struct {
+	address *string
+	key     *string
+	hash    *string
+	cert    *string
+
+	cmd                *string
+	password           *string
+	add                *bool
+	delete             *bool
+	modify             *bool
+	list               *bool
+	uploadSpeedLimit   *int
+	downloadSpeedLimit *int
+	ipLimit            *int
+	ctx                context.Context
+}
+
+func (apiController) Name() string {
+	return "api"
+}
+
+func (o *apiController) listUsers(apiClient service.TrojanServerServiceClient) error {
+	stream, err := apiClient.ListUsers(o.ctx, &service.ListUsersRequest{})
+	if err != nil {
+		return err
+	}
+	defer stream.CloseSend()
+	result := []service.ListUsersResponse{}
+	for {
+		resp, err := stream.Recv()
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+		result = append(result, *resp)
+	}
+	data, err := json.Marshal(result)
+	common.Must(err)
+	fmt.Println(string(data))
+	return nil
+}
+
+func (o *apiController) getUsers(apiClient service.TrojanServerServiceClient) error {
+	stream, err := apiClient.GetUsers(o.ctx)
+	if err != nil {
+		return err
+	}
+	defer stream.CloseSend()
+	err = stream.Send(&service.GetUsersRequest{
+		User: &service.User{
+			Password: *o.password,
+			Hash:     *o.hash,
+		},
+	})
+	if err != nil {
+		return err
+	}
+	resp, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	data, err := json.Marshal(resp)
+	common.Must(err)
+	fmt.Print(string(data))
+	return nil
+}
+
+func (o *apiController) setUsers(apiClient service.TrojanServerServiceClient) error {
+	stream, err := apiClient.SetUsers(o.ctx)
+	if err != nil {
+		return err
+	}
+	defer stream.CloseSend()
+
+	req := &service.SetUsersRequest{
+		User: &service.User{
+			Password: *o.password,
+			Hash:     *o.hash,
+		},
+		IpLimit: int32(*o.ipLimit),
+		SpeedLimit: &service.Speed{
+			UploadSpeed:   uint64(*o.uploadSpeedLimit),
+			DownloadSpeed: uint64(*o.downloadSpeedLimit),
+		},
+	}
+	if *o.add {
+		req.Operation = service.SetUsersRequest_Add
+	} else if *o.modify {
+		req.Operation = service.SetUsersRequest_Modify
+	} else if *o.delete {
+		req.Operation = service.SetUsersRequest_Delete
+	} else {
+		return common.NewError("Invalid operation")
+	}
+
+	err = stream.Send(req)
+	if err != nil {
+		return err
+	}
+	resp, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	if resp.Success {
+		fmt.Println("Done")
+	} else {
+		fmt.Println("Failed: " + resp.Info)
+	}
+	return nil
+}
+
+func (o *apiController) Handle() error {
+	if *o.cmd == "" {
+		return common.NewError("")
+	}
+	conn, err := grpc.Dial(*o.address, grpc.WithInsecure())
+	if err != nil {
+		log.Error(err)
+		return nil
+	}
+	defer conn.Close()
+	apiClient := service.NewTrojanServerServiceClient(conn)
+	switch *o.cmd {
+	case "list":
+		err := o.listUsers(apiClient)
+		if err != nil {
+			log.Error(err)
+		}
+	case "get":
+		err := o.getUsers(apiClient)
+		if err != nil {
+			log.Error(err)
+		}
+	case "set":
+		err := o.setUsers(apiClient)
+		if err != nil {
+			log.Error(err)
+		}
+	default:
+		log.Error("unknown command " + *o.cmd)
+	}
+	return nil
+}
+
+func (o *apiController) Priority() int {
+	return 50
+}
+
+func init() {
+	option.RegisterHandler(&apiController{
+		cmd:                flag.String("api", "", "Connect to a Trojan-Go API service. \"-api add/get/list\""),
+		address:            flag.String("api-addr", "127.0.0.1:10000", "Address of Trojan-Go API service"),
+		password:           flag.String("target-password", "", "Password of the target user"),
+		hash:               flag.String("target-hash", "", "Hash of the target user"),
+		add:                flag.Bool("add-profile", false, "Add a new profile with API"),
+		delete:             flag.Bool("delete-profile", false, "Delete an existing profile with API"),
+		modify:             flag.Bool("modify-profile", false, "Modify an existing profile with API"),
+		uploadSpeedLimit:   flag.Int("upload-speed-limit", 0, "Limit the upload speed with API"),
+		downloadSpeedLimit: flag.Int("download-speed-limit", 0, "Limit the download speed with API"),
+		ipLimit:            flag.Int("ip-limit", 0, "Limit the number of IP with API"),
+		ctx:                context.Background(),
+	})
+}

api/control/control_test.go

@@ -0,0 +1 @@
+package control

api/service/client_test.go

@@ -28,6 +28,8 @@ func TestClientAPI(t *testing.T) {
 	auth, err := memory.NewAuthenticator(ctx)
 	common.Must(err)
 	go RunClientAPI(ctx, auth)
+
+	time.Sleep(time.Second * 3)
 	common.Must(auth.AddUser("hash1234"))
 	valid, user := auth.AuthUser("hash1234")
 	if !valid {

api/service/server_test.go

@@ -28,6 +28,7 @@ func TestServerAPI(t *testing.T) {
 	auth, err := memory.NewAuthenticator(ctx)
 	common.Must(err)
 	go RunServerAPI(ctx, auth)
+	time.Sleep(time.Second * 3)
 	common.Must(auth.AddUser("hash1234"))
 	_, user := auth.AuthUser("hash1234")
 	conn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", port), grpc.WithInsecure())