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  }