sshpoke/internal/server/driver/ssh/driver.go

package ssh

import (
	"context"
	"errors"
	"path"
	"sync"

	"github.com/Neur0toxine/sshpoke/internal/config"
	"github.com/Neur0toxine/sshpoke/internal/server/driver/base"
	"github.com/Neur0toxine/sshpoke/internal/server/driver/ssh/sshproto"
	"github.com/Neur0toxine/sshpoke/internal/server/driver/ssh/types"
	"github.com/Neur0toxine/sshpoke/internal/server/driver/util"
	"github.com/Neur0toxine/sshpoke/pkg/dto"
	"golang.org/x/crypto/ssh"
)

type SSH struct {
	base.Base
	params Params
	proto  *sshproto.Client
	wg     sync.WaitGroup
}

func New(ctx context.Context, name string, params config.DriverParams) (base.Driver, error) {
	drv := &SSH{
		Base: base.New(ctx, name),
	}
	if err := util.UnmarshalParams(params, &drv.params); err != nil {
		return nil, err
	}
	drv.populateFromSSHConfig()
	drv.proto = sshproto.New(drv.params.Address, drv.params.Auth.User, drv.authenticators(), drv.Log())
	go drv.proto.Connect(drv.Context())
	return drv, nil
}

func (d *SSH) populateFromSSHConfig() {
	if d.params.Auth.Directory == "" {
		return
	}
	cfg, err := parseSSHConfig(types.SmartPath(path.Join(string(d.params.Auth.Directory), "config")))
	if err != nil {
		return
	}
	if user, err := cfg.Get(d.params.Address, "User"); err == nil && user != "" {
		d.params.Auth.User = user
	}
	if usePass, err := cfg.Get(d.params.Address, "PasswordAuthentication"); err == nil && usePass == "yes" {
		d.params.Auth.Type = types.AuthTypePassword
	}
	if keyfile, err := cfg.Get(d.params.Address, "IdentityFile"); err == nil && keyfile != "" {
		resolvedKeyFile, err := types.SmartPath(keyfile).Resolve(false)
		if err == nil {
			d.params.Auth.Type = types.AuthTypeKey
			d.params.Auth.Keyfile = resolvedKeyFile
		}
	}
}

func (d *SSH) Handle(event dto.Event) error {
	// TODO: Implement event handling & connections management.
	return errors.New("server handler is not implemented yet")
}

func (d *SSH) Driver() config.DriverType {
	return config.DriverSSH
}

func (d *SSH) WaitForShutdown() {
	d.wg.Wait()
}

func (d *SSH) authenticators() []ssh.AuthMethod {
	auth := d.authenticator()
	if auth == nil {
		return nil
	}
	return []ssh.AuthMethod{auth}
}

func (d *SSH) authenticator() ssh.AuthMethod {
	switch d.params.Auth.Type {
	case types.AuthTypePasswordless:
		return sshproto.AuthPassword("")
	case types.AuthTypePassword:
		return sshproto.AuthPassword(d.params.Auth.Password)
	case types.AuthTypeKey:
		if d.params.Auth.Keyfile != "" {
			keyAuth, err := sshproto.AuthKeyFile(
				types.SmartPath(path.Join(d.params.Auth.Directory.String(), d.params.Auth.Keyfile)))
			if err != nil {
				return nil
			}
			return keyAuth
		}
		dirAuth, err := sshproto.AuthKeyDir(d.params.Auth.Directory)
		if err != nil {
			return nil
		}
		return dirAuth
	}
	return nil
}