hermes/backend/sshappbackend/remote-code/main.go

460 lines
11 KiB
Go

package main
import (
"errors"
"fmt"
"net"
"os"
"strconv"
"strings"
"sync"
"git.terah.dev/imterah/hermes/backend/commonbackend"
"git.terah.dev/imterah/hermes/backend/sshappbackend/datacommands"
"git.terah.dev/imterah/hermes/backend/sshappbackend/remote-code/backendutil_custom"
"github.com/charmbracelet/log"
)
type TCPProxy struct {
connectionIDIndex uint16
connectionIDLock sync.Mutex
proxyInformation *commonbackend.AddProxy
connections map[uint16]net.Conn
server net.Listener
}
type UDPProxy struct {
server *net.UDPConn
proxyInformation *commonbackend.AddProxy
}
type SSHRemoteAppBackend struct {
proxyIDIndex uint16
proxyIDLock sync.Mutex
tcpProxies map[uint16]*TCPProxy
udpProxies map[uint16]*UDPProxy
isRunning bool
sock net.Conn
}
func (backend *SSHRemoteAppBackend) StartBackend(byte []byte) (bool, error) {
backend.tcpProxies = map[uint16]*TCPProxy{}
backend.udpProxies = map[uint16]*UDPProxy{}
backend.isRunning = true
return true, nil
}
func (backend *SSHRemoteAppBackend) StopBackend() (bool, error) {
for tcpProxyIndex, tcpProxy := range backend.tcpProxies {
for _, tcpConnection := range tcpProxy.connections {
tcpConnection.Close()
}
tcpProxy.server.Close()
delete(backend.tcpProxies, tcpProxyIndex)
}
for udpProxyIndex, udpProxy := range backend.udpProxies {
udpProxy.server.Close()
delete(backend.udpProxies, udpProxyIndex)
}
backend.isRunning = false
return true, nil
}
func (backend *SSHRemoteAppBackend) GetBackendStatus() (bool, error) {
return backend.isRunning, nil
}
func (backend *SSHRemoteAppBackend) StartProxy(command *commonbackend.AddProxy) (uint16, bool, error) {
// Allocate a new proxy ID
backend.proxyIDLock.Lock()
proxyID := backend.proxyIDIndex
backend.proxyIDIndex++
backend.proxyIDLock.Unlock()
if command.Protocol == "tcp" {
backend.tcpProxies[proxyID] = &TCPProxy{
connections: map[uint16]net.Conn{},
proxyInformation: command,
}
server, err := net.Listen("tcp", fmt.Sprintf(":%d", command.DestPort))
if err != nil {
return 0, false, fmt.Errorf("failed to open server: %s", err.Error())
}
backend.tcpProxies[proxyID].server = server
go func() {
for {
conn, err := server.Accept()
if err != nil {
log.Warnf("failed to accept connection: %s", err.Error())
return
}
go func() {
backend.tcpProxies[proxyID].connectionIDLock.Lock()
connectionID := backend.tcpProxies[proxyID].connectionIDIndex
backend.tcpProxies[proxyID].connectionIDIndex++
backend.tcpProxies[proxyID].connectionIDLock.Unlock()
backend.tcpProxies[proxyID].connections[connectionID] = conn
dataBuf := make([]byte, 65535)
onConnection := &datacommands.TCPConnectionOpened{
ProxyID: proxyID,
ConnectionID: connectionID,
}
connectionCommandMarshalled, err := datacommands.Marshal(onConnection)
if err != nil {
log.Errorf("failed to marshal connection message: %s", err.Error())
}
backend.sock.Write(connectionCommandMarshalled)
tcpData := &datacommands.TCPProxyData{
ProxyID: proxyID,
ConnectionID: connectionID,
}
for {
len, err := conn.Read(dataBuf)
if err != nil {
if errors.Is(err, net.ErrClosed) {
return
} else if err.Error() != "EOF" {
log.Warnf("failed to read from sock: %s", err.Error())
}
conn.Close()
break
}
tcpData.DataLength = uint16(len)
marshalledMessageCommand, err := datacommands.Marshal(tcpData)
if err != nil {
log.Warnf("failed to marshal message data: %s", err.Error())
conn.Close()
break
}
if _, err := backend.sock.Write(marshalledMessageCommand); err != nil {
log.Warnf("failed to send marshalled message data: %s", err.Error())
conn.Close()
break
}
if _, err := backend.sock.Write(dataBuf[:len]); err != nil {
log.Warnf("failed to send raw message data: %s", err.Error())
conn.Close()
break
}
}
onDisconnect := &datacommands.TCPConnectionClosed{
ProxyID: proxyID,
ConnectionID: connectionID,
}
disconnectionCommandMarshalled, err := datacommands.Marshal(onDisconnect)
if err != nil {
log.Errorf("failed to marshal disconnection message: %s", err.Error())
}
backend.sock.Write(disconnectionCommandMarshalled)
}()
}
}()
} else if command.Protocol == "udp" {
backend.udpProxies[proxyID] = &UDPProxy{
proxyInformation: command,
}
server, err := net.ListenUDP("udp", &net.UDPAddr{
IP: net.IPv4(0, 0, 0, 0),
Port: int(command.DestPort),
})
if err != nil {
return 0, false, fmt.Errorf("failed to open server: %s", err.Error())
}
backend.udpProxies[proxyID].server = server
dataBuf := make([]byte, 65535)
udpProxyData := &datacommands.UDPProxyData{
ProxyID: proxyID,
}
go func() {
for {
len, addr, err := server.ReadFromUDP(dataBuf)
if err != nil {
log.Warnf("failed to read from UDP socket: %s", err.Error())
continue
}
udpProxyData.ClientIP = addr.IP.String()
udpProxyData.ClientPort = uint16(addr.Port)
udpProxyData.DataLength = uint16(len)
marshalledMessageCommand, err := datacommands.Marshal(udpProxyData)
if err != nil {
log.Warnf("failed to marshal message data: %s", err.Error())
continue
}
if _, err := backend.sock.Write(marshalledMessageCommand); err != nil {
log.Warnf("failed to send marshalled message data: %s", err.Error())
continue
}
if _, err := backend.sock.Write(dataBuf[:len]); err != nil {
log.Warnf("failed to send raw message data: %s", err.Error())
continue
}
}
}()
}
return proxyID, true, nil
}
func (backend *SSHRemoteAppBackend) StopProxy(command *datacommands.RemoveProxy) (bool, error) {
tcpProxy, ok := backend.tcpProxies[command.ProxyID]
if !ok {
udpProxy, ok := backend.udpProxies[command.ProxyID]
if !ok {
return ok, fmt.Errorf("could not find proxy")
}
udpProxy.server.Close()
delete(backend.udpProxies, command.ProxyID)
} else {
for _, tcpConnection := range tcpProxy.connections {
tcpConnection.Close()
}
tcpProxy.server.Close()
delete(backend.tcpProxies, command.ProxyID)
}
return true, nil
}
func (backend *SSHRemoteAppBackend) GetAllProxies() []uint16 {
proxyList := make([]uint16, len(backend.tcpProxies)+len(backend.udpProxies))
currentPos := 0
for tcpProxy := range backend.tcpProxies {
proxyList[currentPos] = tcpProxy
currentPos += 1
}
for udpProxy := range backend.udpProxies {
proxyList[currentPos] = udpProxy
currentPos += 1
}
return proxyList
}
func (backend *SSHRemoteAppBackend) ResolveProxy(proxyID uint16) *datacommands.ProxyInformationResponse {
var proxyInformation *commonbackend.AddProxy
response := &datacommands.ProxyInformationResponse{}
tcpProxy, ok := backend.tcpProxies[proxyID]
if !ok {
udpProxy, ok := backend.udpProxies[proxyID]
if !ok {
response.Exists = false
return response
}
proxyInformation = udpProxy.proxyInformation
} else {
proxyInformation = tcpProxy.proxyInformation
}
response.Exists = true
response.SourceIP = proxyInformation.SourceIP
response.SourcePort = proxyInformation.SourcePort
response.DestPort = proxyInformation.DestPort
response.Protocol = proxyInformation.Protocol
return response
}
func (backend *SSHRemoteAppBackend) GetAllClientConnections(proxyID uint16) []uint16 {
tcpProxy, ok := backend.tcpProxies[proxyID]
if !ok {
return []uint16{}
}
connectionsArray := make([]uint16, len(tcpProxy.connections))
currentPos := 0
for connectionIndex := range tcpProxy.connections {
connectionsArray[currentPos] = connectionIndex
currentPos++
}
return connectionsArray
}
func (backend *SSHRemoteAppBackend) ResolveConnection(proxyID, connectionID uint16) *datacommands.ProxyConnectionInformationResponse {
response := &datacommands.ProxyConnectionInformationResponse{}
tcpProxy, ok := backend.tcpProxies[proxyID]
if !ok {
response.Exists = false
return response
}
connection, ok := tcpProxy.connections[connectionID]
if !ok {
response.Exists = false
return response
}
addr := connection.RemoteAddr().String()
ip := addr[:strings.LastIndex(addr, ":")]
port, err := strconv.Atoi(addr[strings.LastIndex(addr, ":")+1:])
if err != nil {
log.Warnf("failed to parse client port: %s", err.Error())
response.Exists = false
return response
}
response.ClientIP = ip
response.ClientPort = uint16(port)
return response
}
func (backend *SSHRemoteAppBackend) CheckParametersForConnections(clientParameters *commonbackend.CheckClientParameters) *commonbackend.CheckParametersResponse {
return &commonbackend.CheckParametersResponse{
IsValid: true,
}
}
func (backend *SSHRemoteAppBackend) CheckParametersForBackend(arguments []byte) *commonbackend.CheckParametersResponse {
return &commonbackend.CheckParametersResponse{
IsValid: true,
}
}
func (backend *SSHRemoteAppBackend) HandleTCPMessage(message *datacommands.TCPProxyData, data []byte) {
tcpProxy, ok := backend.tcpProxies[message.ProxyID]
if !ok {
log.Warnf("could not find tcp proxy (ID %d)", message.ProxyID)
return
}
connection, ok := tcpProxy.connections[message.ConnectionID]
if !ok {
log.Warnf("could not find tcp proxy (ID %d) with connection ID (%d)", message.ProxyID, message.ConnectionID)
return
}
connection.Write(data)
}
func (backend *SSHRemoteAppBackend) HandleUDPMessage(message *datacommands.UDPProxyData, data []byte) {
udpProxy, ok := backend.udpProxies[message.ProxyID]
if !ok {
return
}
udpProxy.server.WriteToUDP(data, &net.UDPAddr{
IP: net.ParseIP(message.ClientIP),
Port: int(message.ClientPort),
})
}
func (backend *SSHRemoteAppBackend) OnTCPConnectionClosed(proxyID, connectionID uint16) {
tcpProxy, ok := backend.tcpProxies[proxyID]
if !ok {
return
}
connection, ok := tcpProxy.connections[connectionID]
if !ok {
return
}
connection.Close()
delete(tcpProxy.connections, connectionID)
}
func (backend *SSHRemoteAppBackend) OnSocketConnection(sock net.Conn) {
backend.sock = sock
}
func main() {
logLevel := os.Getenv("HERMES_LOG_LEVEL")
if logLevel != "" {
switch logLevel {
case "debug":
log.SetLevel(log.DebugLevel)
case "info":
log.SetLevel(log.InfoLevel)
case "warn":
log.SetLevel(log.WarnLevel)
case "error":
log.SetLevel(log.ErrorLevel)
case "fatal":
log.SetLevel(log.FatalLevel)
}
}
backend := &SSHRemoteAppBackend{}
application := backendutil_custom.NewHelper(backend)
err := application.Start()
if err != nil {
log.Fatalf("failed execution in application: %s", err.Error())
}
}