package sshtun

import (
	"bufio"
	"io"

	"github.com/Neur0toxine/sshpoke/pkg/proto/ssh"
	"github.com/jonstacks/iomerge"
	"go.uber.org/zap"
)

func OutputReaderCallback(callback func(string)) SessionCallback {
	return func(session *ssh.Session) {
		stdout, err := session.StdoutPipe()
		if err != nil {
			return
		}
		stderr, err := session.StderrPipe()
		if err != nil {
			return
		}
		combined := iomerge.NewStreamMerger(2)
		combined.Write(func(readers chan<- io.Reader) {
			readers <- stdout
			readers <- stderr
		})
		combined.Read(func(r io.Reader) error {
			scan := bufio.NewScanner(r)
			for scan.Scan() {
				callback(scan.Text())
			}
			return nil
		})
		_ = combined.Wait()
	}
}

func BannerDebugLogCallback(log *zap.SugaredLogger) ssh.BannerCallback {
	return func(msg string) error {
		log.Debug("banner: ", msg)
		return nil
	}
}