Merge pull request #9 from zxyxx/master

Add BlockSameIP

lib/cli/dispatcher/block_same_ip.go

@@ -0,0 +1,41 @@
+package dispatcher
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/WangYihang/Platypus/lib/context"
+	"github.com/WangYihang/Platypus/lib/util/log"
+)
+
+func (dispatcher Dispatcher) BlockSameIP(args []string) {
+	if len(args) != 1 {
+		log.Error("Arguments error, use `Help BlockSameIP` to get more information")
+		dispatcher.BlockSameIPHelp([]string{})
+		return
+	}
+	parseInt,err := strconv.Atoi(args[0])
+	if err != nil {
+		log.Error("Something error")
+		return
+	}
+	if parseInt == 1 {
+		log.Success("BlockSameIP set to 1, will only accept one client from every unique IP")
+		context.Ctx.BlockSameIP = 1
+	} else if parseInt == 0 {
+		log.Success("BlockSameIP set to 0, every IP can have many clients")
+		context.Ctx.BlockSameIP = 0
+	}
+}
+
+func (dispatcher Dispatcher) BlockSameIPHelp(args []string) {
+	fmt.Println("Usage of BlockSameIP")
+	fmt.Println("\tBlockSameIP [01]")
+	fmt.Println("\tWhen BlockSameIP set to 1, will only accept one client from every unique IP, by default")
+	fmt.Println("\tWhen BlockSameIP set to 0, every IP can have many clients")
+}
+
+func (dispatcher Dispatcher) BlockSameIPDesc(args []string) {
+	fmt.Println("BlockSameIP")
+	fmt.Println("\tIf a client is online, decline other requests from the same IP")
+}

lib/cli/dispatcher/jump.go

@@ -10,7 +10,7 @@ import (
 
 func (dispatcher Dispatcher) Jump(args []string) {
 	if len(args) != 1 {
-		log.Error("Arguments error, use `Help Jump` to get more Jumprmation")
+		log.Error("Arguments error, use `Help Jump` to get more information")
 		dispatcher.JumpHelp([]string{})
 		return
 	}

lib/context/context.go

@@ -4,6 +4,7 @@ type Context struct {
 	Servers       map[string](*TCPServer)
 	Current       *TCPClient
 	CommandPrompt string
+	BlockSameIP int
 }
 
 var Ctx *Context
@@ -14,6 +15,7 @@ func CreateContext() {
 			Servers:       make(map[string](*TCPServer)),
 			Current:       nil,
 			CommandPrompt: ">> ",
+			BlockSameIP:   1,
 		}
 	}
 }

lib/context/server.go

@@ -80,7 +80,6 @@ func (s *TCPServer) Run() {
 			continue
 		}
 		client := CreateTCPClient(conn)
-		log.Info("New client %s Connected", client.Desc())
 		// Reverse shell as a service
 		buffer := make([]byte, 4)
 		client.Conn.SetReadDeadline(time.Now().Add(time.Second * 3))
@@ -141,7 +140,26 @@ func (s *TCPServer) Run() {
 			Ctx.DeleteTCPClient(client)
 			log.Info("RaaS: %s", command)
 		} else {
-			s.AddTCPClient(client)
+			switch Ctx.BlockSameIP {
+			case 1:
+				newclientIP := client.Conn.RemoteAddr().String()
+				newclientIP = strings.Split(newclientIP, ":")[0]
+				clientExist := 0
+				for _, client := range s.Clients {
+					clientIP := client.Conn.RemoteAddr().String()
+					clientIP = strings.Split(clientIP, ":")[0]
+					if newclientIP == clientIP {
+						clientExist = 1
+					}
+				}
+				if clientExist == 0 {
+					log.Info("New client %s Connected", client.Desc())
+					s.AddTCPClient(client)
+				}
+			case 0:
+				log.Info("New client %s Connected", client.Desc())
+				s.AddTCPClient(client)
+			}
 		}
 	}
 }