Files
oneterm/backend/pkg/proto/ssh/handler/handlers.go
2024-02-01 20:53:29 +08:00

127 lines
2.9 KiB
Go

package handler
import (
"io"
"os"
"strings"
"time"
gossh "github.com/gliderlabs/ssh"
gssh "golang.org/x/crypto/ssh"
"github.com/veops/oneterm/pkg/logger"
"github.com/veops/oneterm/pkg/proto/ssh/api"
cfg "github.com/veops/oneterm/pkg/proto/ssh/config"
"github.com/veops/oneterm/pkg/server/model"
)
type sshdServer struct {
Core *api.CoreInstance
}
func Init(address, apiHost, token, privateKeyPath, secretKey string) (*gossh.Server, error) {
sshd := NewSshdServer(apiHost, token, secretKey)
s := &gossh.Server{
Addr: address,
Handler: sshd.HomeHandler,
PasswordHandler: sshd.PasswordHandler,
PublicKeyHandler: sshd.PublicKeyHandler,
IdleTimeout: time.Hour*2 + time.Minute,
}
for _, v := range hostPrivateKeys(privateKeyPath) {
singer, er := gssh.ParsePrivateKey(v)
if er != nil {
continue
}
s.AddHostKey(singer)
}
return s, nil
}
func hostPrivateKeys(privateKeyPath string) [][]byte {
var res [][]byte
if privateKeyPath == "" {
homeDir, er := os.UserHomeDir()
if er != nil {
logger.L.Error(er.Error())
}
privateKeyPath = homeDir + "/.ssh/id_ed25519"
}
privateKey, err := os.ReadFile(privateKeyPath)
if err != nil {
logger.L.Error(err.Error())
return res
}
return [][]byte{privateKey}
}
func NewSshdServer(apiHost, token, secretKey string) *sshdServer {
s := &sshdServer{
Core: api.NewCoreInstance(apiHost, token, secretKey),
}
return s
}
func (s *sshdServer) PasswordHandler(ctx gossh.Context, password string) bool {
if password == "" {
return false
}
if ctx.User() == cfg.SSHConfig.WebUser && password == cfg.SSHConfig.WebToken {
ctx.SetValue("sshType", model.SESSIONTYPE_WEB)
return true
}
ctx.SetValue("sshType", model.SESSIONTYPE_CLIENT)
s.Core.Auth.Username = ctx.User()
s.Core.Auth.Password = password
s.Core.Auth.PublicKey = ""
return s.Auth(ctx)
}
func (s *sshdServer) PublicKeyHandler(ctx gossh.Context, key gossh.PublicKey) bool {
authorizedKey := gssh.MarshalAuthorizedKey(key)
s.Core.Auth.PublicKey = strings.TrimSpace(string(authorizedKey))
if s.Core.Auth.PublicKey == "" {
return false
}
s.Core.Auth.Username = ctx.User()
s.Core.Auth.Password = ""
if ctx.Value("sshType") == nil {
ctx.SetValue("sshType", model.SESSIONTYPE_CLIENT)
}
return s.Auth(ctx)
}
func (s *sshdServer) Auth(ctx gossh.Context) bool {
cookie, err := s.Core.Auth.Authenticate()
if err != nil || cookie == "" {
return false
}
ctx.SetValue("cookie", cookie)
return true
}
func (s *sshdServer) HomeHandler(gs gossh.Session) {
if py, winChan, isPty := gs.Pty(); isPty {
if py.Window.Height == 0 {
py.Window.Height = 24
}
interactiveSrv := NewInteractiveHandler(gs, s, py)
go interactiveSrv.WatchWinSize(winChan)
interactiveSrv.Schedule(&py)
} else {
if _, err := io.WriteString(gs, "不是PTY请求.\n"); err != nil {
logger.L.Error(err.Error())
}
err := gs.Exit(1)
if err != nil {
logger.L.Error(err.Error())
}
return
}
}