You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
286 lines
6.9 KiB
Go
286 lines
6.9 KiB
Go
package net
|
|
|
|
import (
|
|
"b612.me/stario"
|
|
"b612.me/starlog"
|
|
"context"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type TcpConn struct {
|
|
*net.TCPConn
|
|
f *os.File
|
|
}
|
|
|
|
type TcpServer struct {
|
|
LocalAddr string
|
|
UsingKeepAlive bool
|
|
KeepAlivePeriod int
|
|
KeepAliveIdel int
|
|
KeepAliveCount int
|
|
sync.Mutex
|
|
Clients map[string]*TcpConn
|
|
Interactive bool
|
|
UserTimeout int
|
|
ShowRecv bool
|
|
ShowAsHex bool
|
|
SaveToFolder string
|
|
Listen *net.TCPListener
|
|
LogPath string
|
|
stopCtx context.Context
|
|
stopFn context.CancelFunc
|
|
}
|
|
|
|
func (s *TcpServer) Close() error {
|
|
return s.Listen.Close()
|
|
}
|
|
|
|
func (s *TcpServer) handleInteractive() {
|
|
var conn *TcpConn
|
|
var currentCmd string
|
|
notifyMap := make(map[string]chan struct{})
|
|
if !s.Interactive {
|
|
return
|
|
}
|
|
starlog.Infoln("Interactive mode enabled")
|
|
for {
|
|
select {
|
|
case <-s.stopCtx.Done():
|
|
starlog.Infoln("Interactive mode stopped due to context done")
|
|
return
|
|
default:
|
|
}
|
|
cmd := stario.MessageBox("", "").MustString()
|
|
if cmd == "" {
|
|
continue
|
|
}
|
|
cmdf := strings.Fields(cmd)
|
|
switch cmdf[0] {
|
|
case "list":
|
|
s.Lock()
|
|
for k, v := range s.Clients {
|
|
starlog.Green("Client %s: %s\n", k, v.RemoteAddr().String())
|
|
}
|
|
s.Unlock()
|
|
case "use":
|
|
if len(cmdf) < 2 {
|
|
starlog.Errorln("use command need a client address")
|
|
continue
|
|
}
|
|
conn = s.Clients[cmdf[1]]
|
|
if conn == nil {
|
|
starlog.Errorln("Client not found")
|
|
continue
|
|
}
|
|
starlog.Infof("Using client %s\n", conn.RemoteAddr().String())
|
|
case "hex":
|
|
currentCmd = "hex"
|
|
starlog.Infoln("Switch to hex mode,send hex to remote client")
|
|
case "text":
|
|
currentCmd = "text"
|
|
starlog.Infoln("Switch to text mode,send text to remote client")
|
|
case "close":
|
|
if conn.TCPConn == nil {
|
|
starlog.Errorln("No client selected")
|
|
continue
|
|
}
|
|
conn.TCPConn.Close()
|
|
starlog.Infof("Client %s closed\n", conn.RemoteAddr().String())
|
|
conn = nil
|
|
currentCmd = ""
|
|
case "startauto":
|
|
if conn == nil {
|
|
starlog.Errorln("No client selected")
|
|
continue
|
|
}
|
|
notifyMap[conn.RemoteAddr().String()] = make(chan struct{})
|
|
go func(conn *TcpConn) {
|
|
for {
|
|
select {
|
|
case <-notifyMap[conn.RemoteAddr().String()]:
|
|
starlog.Infoln("Auto send stopped")
|
|
return
|
|
default:
|
|
}
|
|
_, err := conn.Write([]byte(strings.Repeat("B612", 256)))
|
|
if err != nil {
|
|
starlog.Errorln("Write error:", err)
|
|
return
|
|
}
|
|
}
|
|
}(conn)
|
|
starlog.Infoln("Auto send started")
|
|
case "closeauto":
|
|
if conn == nil {
|
|
starlog.Errorln("No client selected")
|
|
continue
|
|
}
|
|
close(notifyMap[conn.RemoteAddr().String()])
|
|
case "send":
|
|
if conn == nil {
|
|
starlog.Errorln("No client selected")
|
|
continue
|
|
}
|
|
if currentCmd == "hex" {
|
|
data, err := hex.DecodeString(strings.TrimSpace(strings.TrimPrefix(cmd, "send")))
|
|
if err != nil {
|
|
starlog.Errorln("Hex decode error:", err)
|
|
continue
|
|
}
|
|
_, err = conn.Write(data)
|
|
if err != nil {
|
|
starlog.Errorln("Write error:", err)
|
|
} else {
|
|
if conn.f != nil {
|
|
conn.f.Write([]byte(time.Now().String() + " send\n"))
|
|
conn.f.Write(data)
|
|
conn.f.Write([]byte("\n"))
|
|
}
|
|
}
|
|
} else {
|
|
_, err := conn.Write([]byte(strings.TrimSpace(strings.TrimPrefix(cmd, "send"))))
|
|
if err != nil {
|
|
starlog.Errorln("Write error:", err)
|
|
} else {
|
|
if conn.f != nil {
|
|
conn.f.Write([]byte(time.Now().String() + " send\n"))
|
|
conn.f.Write([]byte(cmdf[1]))
|
|
conn.f.Write([]byte("\n"))
|
|
}
|
|
}
|
|
}
|
|
starlog.Infof("Send to %s success\n", conn.RemoteAddr().String())
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *TcpServer) Run() error {
|
|
s.stopCtx, s.stopFn = context.WithCancel(context.Background())
|
|
if s.LogPath != "" {
|
|
err := starlog.SetLogFile(s.LogPath, starlog.Std, true)
|
|
if err != nil {
|
|
starlog.Errorln("SetLogFile error:", err)
|
|
return fmt.Errorf("SetLogFile error: %w", err)
|
|
}
|
|
}
|
|
s.Clients = make(map[string]*TcpConn)
|
|
tcpAddr, err := net.ResolveTCPAddr("tcp", s.LocalAddr)
|
|
if err != nil {
|
|
starlog.Errorln("ResolveTCPAddr error:", err)
|
|
return fmt.Errorf("ResolveTCPAddr error: %w", err)
|
|
}
|
|
s.Listen, err = net.ListenTCP("tcp", tcpAddr)
|
|
if err != nil {
|
|
starlog.Errorln("ListenTCP error:", err)
|
|
return fmt.Errorf("ListenTCP error: %w", err)
|
|
}
|
|
starlog.Infof("TcpServer listen on %s\n", s.LocalAddr)
|
|
if s.Interactive {
|
|
go s.handleInteractive()
|
|
}
|
|
for {
|
|
select {
|
|
case <-s.stopCtx.Done():
|
|
starlog.Infoln("TcpServer stopped due to context done")
|
|
return s.Listen.Close()
|
|
default:
|
|
}
|
|
conn, err := s.Listen.AcceptTCP()
|
|
if err != nil {
|
|
starlog.Errorln("AcceptTCP error:", err)
|
|
continue
|
|
}
|
|
starlog.Infof("Accept new connection from %s\n", conn.RemoteAddr().String())
|
|
s.Lock()
|
|
s.Clients[conn.RemoteAddr().String()] = s.getTcpConn(conn)
|
|
s.Unlock()
|
|
go s.handleConn(s.Clients[conn.RemoteAddr().String()])
|
|
}
|
|
}
|
|
|
|
func (s *TcpServer) getTcpConn(conn *net.TCPConn) *TcpConn {
|
|
var err error
|
|
var f *os.File
|
|
if s.SaveToFolder != "" {
|
|
f, err = os.Create(filepath.Join(s.SaveToFolder, strings.ReplaceAll(conn.RemoteAddr().String(), ":", "_")))
|
|
if err != nil {
|
|
starlog.Errorf("Create file error for %s: %v\n", conn.RemoteAddr().String(), err)
|
|
}
|
|
}
|
|
return &TcpConn{
|
|
TCPConn: conn,
|
|
f: f,
|
|
}
|
|
}
|
|
|
|
func (s *TcpServer) handleConn(conn *TcpConn) {
|
|
var err error
|
|
log := starlog.Std.NewFlag()
|
|
err = SetTcpInfo(conn.TCPConn, s.UsingKeepAlive, s.KeepAliveIdel, s.KeepAlivePeriod, s.KeepAliveCount, s.UserTimeout)
|
|
if err != nil {
|
|
log.Errorf("SetTcpInfo error for %s: %v\n", conn.RemoteAddr().String(), err)
|
|
s.Lock()
|
|
delete(s.Clients, conn.RemoteAddr().String())
|
|
s.Unlock()
|
|
conn.Close()
|
|
return
|
|
}
|
|
log.Infof("SetKeepAlive success for %s\n", conn.RemoteAddr().String())
|
|
log.Infof("KeepAlivePeriod: %d, KeepAliveIdel: %d, KeepAliveCount: %d, UserTimeout: %d\n", s.KeepAlivePeriod, s.KeepAliveIdel, s.KeepAliveCount, s.UserTimeout)
|
|
if runtime.GOOS != "linux" {
|
|
log.Warningln("keepAliveCount and userTimeout only work on linux")
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case <-s.stopCtx.Done():
|
|
log.Infof("Connection from %s closed due to context done\n", conn.RemoteAddr().String())
|
|
s.Lock()
|
|
delete(s.Clients, conn.RemoteAddr().String())
|
|
s.Unlock()
|
|
conn.Close()
|
|
return
|
|
default:
|
|
}
|
|
buf := make([]byte, 8192)
|
|
n, err := conn.Read(buf)
|
|
if err != nil {
|
|
log.Errorf("Read error for %s: %v\n", conn.RemoteAddr().String(), err)
|
|
s.Lock()
|
|
delete(s.Clients, conn.RemoteAddr().String())
|
|
s.Unlock()
|
|
conn.Close()
|
|
return
|
|
}
|
|
if n > 0 {
|
|
if s.ShowRecv {
|
|
if s.ShowAsHex {
|
|
log.Printf("Recv from %s: %x\n", conn.RemoteAddr().String(), buf[:n])
|
|
} else {
|
|
log.Printf("Recv from %s: %s\n", conn.RemoteAddr().String(), string(buf[:n]))
|
|
}
|
|
}
|
|
if conn.f != nil {
|
|
conn.f.Write([]byte(time.Now().String() + " recv\n"))
|
|
conn.f.Write(buf[:n])
|
|
conn.f.Write([]byte("\n"))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *TcpServer) Stop() {
|
|
s.stopFn()
|
|
if s.Listen != nil {
|
|
s.Close()
|
|
}
|
|
}
|