gpu-compute-chain/x/colinearcore/audit/sshconn.go

110 lines
2.3 KiB
Go

package audit
import (
"errors"
"fmt"
"log"
"strconv"
"time"
"github.com/dgraph-io/badger/v2"
"golang.org/x/crypto/ssh"
)
// Maps SSH ping result channel to address that is being pinged
var pingDB *badger.DB
func MountPingDB() {
var err error
pingDB, err = badger.Open(badger.DefaultOptions("").WithInMemory(true))
if err != nil {
// must force crash here, since ping db is absolutely required
log.Fatalf("Failed to mount in-memory db: %s", err)
}
}
// Begin ping to check if SSH is running on an open port on a host.
//
// NOTE: host should include a port number.
func BeginSshPing(host string) error {
k := []byte(host)
err := pingDB.Update(func(txn *badger.Txn) error {
_, err := txn.Get(k)
if err != nil {
if !errors.Is(err, badger.ErrKeyNotFound) {
// key-not-found condition required to write
return err
} else {
goto writePing
}
}
return fmt.Errorf("already pinging host %s", host)
writePing:
if err := txn.Set(k, []byte(strconv.FormatBool(false))); err != nil {
return err
}
return nil
})
if err != nil {
return err
}
sshConfig := &ssh.ClientConfig{
User: "user",
Auth: []ssh.AuthMethod{ssh.Password("pass")},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
Timeout: 100 * time.Second,
}
go dialSshAsync(host, sshConfig)
return nil
}
// Check on ongoing SSH ping
func CheckSshPing(host string) (bool, error) {
k := []byte(host)
var pingStatus bool
err := pingDB.View(func(txn *badger.Txn) error {
resVal, err := txn.Get(k)
if err != nil {
if !errors.Is(err, badger.ErrKeyNotFound) {
return err
} else {
return fmt.Errorf("not pinging host %s", host)
}
}
err = resVal.Value(func(val []byte) error {
pingStatus, err = strconv.ParseBool(string(val))
if err != nil {
return fmt.Errorf("failed to decode ping status for host %s: %s", host, string(val))
}
return nil
})
return err
})
if err != nil {
return false, err
}
return pingStatus, nil
}
func dialSshAsync(host string, conf *ssh.ClientConfig) {
ssh.Dial("tcp", host, conf)
pingDB.Update(func(txn *badger.Txn) error {
k := []byte(host)
return txn.Set(k, []byte(strconv.FormatBool(true)))
})
// any potential resulting error is swallowed since this is run asynchronously
}