From ba57ccdd4531c69bf347e177aebd270fa1cd5672 Mon Sep 17 00:00:00 2001
From: Hossin Asaadi <hossin.asaadi@gmail.com>
Date: Sun, 31 Mar 2024 08:29:15 +0400
Subject: [PATCH] API: add Source IP Block command (#3211)

* add SourceIpBlock to API

* improve inboundTag
---
 main/commands/all/api/api.go             |   1 +
 main/commands/all/api/source_ip_block.go | 132 +++++++++++++++++++++++
 2 files changed, 133 insertions(+)
 create mode 100644 main/commands/all/api/source_ip_block.go

diff --git a/main/commands/all/api/api.go b/main/commands/all/api/api.go
index f408fd90..ea43d107 100644
--- a/main/commands/all/api/api.go
+++ b/main/commands/all/api/api.go
@@ -23,5 +23,6 @@ var CmdAPI = &base.Command{
 		cmdRemoveOutbounds,
 		cmdAddRules,
 		cmdRemoveRules,
+		cmdSourceIpBlock,
 	},
 }
diff --git a/main/commands/all/api/source_ip_block.go b/main/commands/all/api/source_ip_block.go
new file mode 100644
index 00000000..1f773187
--- /dev/null
+++ b/main/commands/all/api/source_ip_block.go
@@ -0,0 +1,132 @@
+package api
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	routerService "github.com/xtls/xray-core/app/router/command"
+	cserial "github.com/xtls/xray-core/common/serial"
+	"github.com/xtls/xray-core/infra/conf/serial"
+	"github.com/xtls/xray-core/main/commands/base"
+)
+
+var cmdSourceIpBlock = &base.Command{
+	CustomFlags: true,
+	UsageLine:   "{{.Exec}} api sib [--server=127.0.0.1:8080] -outbound=blocked -inbound=socks 1.2.3.4",
+	Short:       "Drop connections by source ip",
+	Long: `
+Drop connections by source ip.
+Arguments:
+	-s, -server 
+		The API server address. Default 127.0.0.1:8080
+	-t, -timeout
+		Timeout seconds to call API. Default 3
+	-outbound
+		route traffic to specific outbound.
+	-inbound
+		target traffig from specific inbound.
+	-ruletag
+		set ruleTag. Default sourceIpBlock
+	-reset
+		remove ruletag and apply new source IPs. Default false
+
+	Example:
+    {{.Exec}} {{.LongName}} --server=127.0.0.1:8080 c1.json c2.json
+`,
+	Run: executeSourceIpBlock,
+}
+
+func executeSourceIpBlock(cmd *base.Command, args []string) {
+	var (
+		inbound  string
+		outbound string
+		ruletag  string
+		reset    bool
+	)
+	setSharedFlags(cmd)
+	cmd.Flag.StringVar(&inbound, "inbound", "", "")
+	cmd.Flag.StringVar(&outbound, "outbound", "", "")
+	cmd.Flag.StringVar(&ruletag, "ruletag", "sourceIpBlock", "")
+	cmd.Flag.BoolVar(&reset, "reset", false, "")
+
+	cmd.Flag.Parse(args)
+
+	unnamedArgs := cmd.Flag.Args()
+	if len(unnamedArgs) == 0 {
+		fmt.Println("reading from stdin:")
+		unnamedArgs = []string{"stdin:"}
+	}
+	conn, ctx, close := dialAPIServer()
+	defer close()
+
+	client := routerService.NewRoutingServiceClient(conn)
+
+	jsonIps, err := json.Marshal(unnamedArgs)
+	if err != nil {
+		fmt.Println("Error marshaling JSON:", err)
+		return
+	}
+
+	jsonInbound, err := json.Marshal([]string{inbound})
+	if inbound == "" {
+		jsonInbound, err = json.Marshal([]string{})
+	}
+	if err != nil {
+		fmt.Println("Error marshaling JSON:", err)
+		return
+	}
+	stringConfig := fmt.Sprintf(`
+	{
+		"routing": {
+			"rules": [
+			  {
+				"ruleTag" : "%s",
+				"inboundTag": %s,		
+				"outboundTag": "%s",
+				"type": "field",
+				"source": %s
+			  }
+			]
+		  }
+	  }
+	  
+	`, ruletag, string(jsonInbound), outbound, string(jsonIps))
+
+	conf, err := serial.DecodeJSONConfig(strings.NewReader(stringConfig))
+	if err != nil {
+		base.Fatalf("failed to decode : %s", err)
+	}
+	rc := *conf.RouterConfig
+
+	config, err := rc.Build()
+	if err != nil {
+		base.Fatalf("failed to build conf: %s", err)
+	}
+	tmsg := cserial.ToTypedMessage(config)
+	if tmsg == nil {
+		base.Fatalf("failed to format config to TypedMessage.")
+	}
+
+	if reset {
+		rr := &routerService.RemoveRuleRequest{
+			RuleTag: ruletag,
+		}
+		resp, err := client.RemoveRule(ctx, rr)
+		if err != nil {
+			base.Fatalf("failed to perform RemoveRule: %s", err)
+		}
+		showJSONResponse(resp)
+
+	}
+	ra := &routerService.AddRuleRequest{
+		Config:       tmsg,
+		ShouldAppend: true,
+	}
+	resp, err := client.AddRule(ctx, ra)
+	if err != nil {
+		base.Fatalf("failed to perform AddRule: %s", err)
+	}
+	showJSONResponse(resp)
+
+}