Skip to content

Commit

Permalink
[+] #28 Support local port frowarding
Browse files Browse the repository at this point in the history
  • Loading branch information
@WangYihang
WangYihang committed Jul 11, 2021
1 parent 6285bfb commit 7591136
Show file tree
Hide file tree
Showing 7 changed files with 317 additions and 167 deletions.
52 changes: 44 additions & 8 deletions lib/cli/dispatcher/tunnel.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,27 +16,63 @@ 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])
dispatcher.TunnelHelp([]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)
}
}

Expand All @@ -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) {
Expand Down
1 change: 0 additions & 1 deletion lib/context/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

Expand Down
106 changes: 106 additions & 0 deletions lib/context/context.go
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,33 @@
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"
"github.com/gin-gonic/gin"
"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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
// }
// }
85 changes: 52 additions & 33 deletions lib/context/server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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")
}

}
}
}
Expand Down
Loading

0 comments on commit 7591136

Please sign in to comment.