mirror of
https://github.com/lwch/natpass
synced 2025-10-06 05:46:49 +08:00
145 lines
3.2 KiB
Go
145 lines
3.2 KiB
Go
package shell
|
|
|
|
import (
|
|
"io"
|
|
"natpass/code/client/pool"
|
|
"natpass/code/network"
|
|
"natpass/code/utils"
|
|
"os"
|
|
|
|
"github.com/lwch/logging"
|
|
"golang.org/x/text/encoding/simplifiedchinese"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
// Link shell link
|
|
type Link struct {
|
|
parent *Shell
|
|
id string // link id
|
|
target string // target id
|
|
targetIdx uint32 // target idx
|
|
remote *pool.Conn
|
|
// in remote
|
|
pid int
|
|
stdin io.WriteCloser
|
|
stdout io.ReadCloser
|
|
// runtime
|
|
sendBytes uint64
|
|
recvBytes uint64
|
|
sendPacket uint64
|
|
recvPacket uint64
|
|
}
|
|
|
|
// GetID get link id
|
|
func (link *Link) GetID() string {
|
|
return link.id
|
|
}
|
|
|
|
// GetBytes get send and recv bytes
|
|
func (link *Link) GetBytes() (uint64, uint64) {
|
|
return link.recvBytes, link.sendBytes
|
|
}
|
|
|
|
// GetPackets get send and recv packets
|
|
func (link *Link) GetPackets() (uint64, uint64) {
|
|
return link.recvPacket, link.sendPacket
|
|
}
|
|
|
|
// SetTargetIdx set link remote index
|
|
func (link *Link) SetTargetIdx(idx uint32) {
|
|
link.targetIdx = idx
|
|
}
|
|
|
|
// Close close link
|
|
func (link *Link) Close() {
|
|
link.onClose()
|
|
p, err := os.FindProcess(link.pid)
|
|
if err == nil {
|
|
p.Kill()
|
|
}
|
|
link.remote.SendDisconnect(link.target, link.targetIdx, link.id)
|
|
link.parent.remove(link.id)
|
|
}
|
|
|
|
// Forward forward data
|
|
func (link *Link) Forward() {
|
|
go link.remoteRead()
|
|
go link.localRead()
|
|
}
|
|
|
|
func (link *Link) remoteRead() {
|
|
defer utils.Recover("remoteRead")
|
|
defer link.Close()
|
|
ch := link.remote.ChanRead(link.id)
|
|
for {
|
|
msg := <-ch
|
|
if msg == nil {
|
|
return
|
|
}
|
|
data, _ := proto.Marshal(msg)
|
|
link.recvBytes += uint64(len(data))
|
|
link.recvPacket++
|
|
link.targetIdx = msg.GetFromIdx()
|
|
switch msg.GetXType() {
|
|
case network.Msg_shell_resize:
|
|
size := msg.GetSresize()
|
|
link.resize(size.GetRows(), size.GetCols())
|
|
case network.Msg_shell_data:
|
|
_, err := link.stdin.Write(msg.GetSdata().GetData())
|
|
if err != nil {
|
|
logging.Error("write data on shell %s link %s failed, err=%v",
|
|
link.parent.Name, link.id, err)
|
|
return
|
|
}
|
|
case network.Msg_disconnect:
|
|
logging.Info("shell %s link %s closed by remote", link.parent.Name, link.id)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (link *Link) localRead() {
|
|
defer utils.Recover("localRead")
|
|
defer link.Close()
|
|
buf := make([]byte, 16*1024)
|
|
for {
|
|
n, err := link.stdout.Read(buf)
|
|
if err != nil {
|
|
logging.Error("read data on shell %s link %s failed, err=%v",
|
|
link.parent.Name, link.id, err)
|
|
return
|
|
}
|
|
if n == 0 {
|
|
continue
|
|
}
|
|
var data []byte
|
|
switch {
|
|
case isUtf8(buf[:n]):
|
|
data = buf[:n]
|
|
case isGBK(buf[:n]):
|
|
data, err = simplifiedchinese.GBK.NewDecoder().Bytes(buf[:n])
|
|
if err != nil {
|
|
logging.Error("transform gbk to utf8 failed: %v", err)
|
|
continue
|
|
}
|
|
}
|
|
logging.Debug("link %s on shell %s read from local %d bytes",
|
|
link.id, link.parent.Name, n)
|
|
send := link.remote.SendShellData(link.target, link.targetIdx, link.id, data)
|
|
link.sendBytes += send
|
|
link.sendPacket++
|
|
}
|
|
}
|
|
|
|
// SendData send data
|
|
func (link *Link) SendData(data []byte) {
|
|
send := link.remote.SendShellData(link.target, link.targetIdx, link.id, data)
|
|
link.sendBytes += send
|
|
link.sendPacket++
|
|
}
|
|
|
|
// SendResize send resize message
|
|
func (link *Link) SendResize(rows, cols uint32) {
|
|
link.remote.SendShellResize(link.target, link.targetIdx, link.id, rows, cols)
|
|
}
|