refactor:Added Integration test and tiny refactoring (#23)

* Refactoring name convention * Added integration test * Added Makefile * Bump golang.org/x/crypto from 0.0.0-20220826181053-bd7e27e6170d to 0.6.0 Bumps [golang.org/x/crypto](https://github.com/golang/crypto) from 0.0.0-20220826181053-bd7e27e6170d to 0.6.0. - [Release notes](https://github.com/golang/crypto/releases) - [Commits](https://github.com/golang/crypto/commits/v0.6.0) --- updated-dependencies: - dependency-name: golang.org/x/crypto dependency-type: direct:production update-type: version-update:semver-minor ... * Upgrade go from 1.16 to 1.20 * Added integration test: HTTP, TCP, SSH * Added Makefile Improve README.md * Fixed unit test CI * Fixed go-version * Added integration test into C.I. actions --------- Signed-off-by: Mario Candela <m4r10.php@gmail.com> Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2025-07-01 18:47:26 -04:00 · 2023-02-26 18:04:05 +01:00
parent bbcc8c9094
commit 6468b5aa6f
23 changed files with 282 additions and 40 deletions
--- a/protocols/strategies/secure_shell.go
+++ b/protocols/strategies/secure_shell.go
@ -0,0 +1,129 @@
+package strategies
+
+import (
+	"beelzebub/parser"
+	"beelzebub/plugins"
+	"beelzebub/tracer"
+	"fmt"
+	"github.com/gliderlabs/ssh"
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh/terminal"
+	"regexp"
+	"strings"
+	"time"
+)
+
+type SecureShellStrategy struct {
+}
+
+func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error {
+	go func() {
+		server := &ssh.Server{
+			Addr:        beelzebubServiceConfiguration.Address,
+			MaxTimeout:  time.Duration(beelzebubServiceConfiguration.DeadlineTimeoutSeconds) * time.Second,
+			IdleTimeout: time.Duration(beelzebubServiceConfiguration.DeadlineTimeoutSeconds) * time.Second,
+			Version:     beelzebubServiceConfiguration.ServerVersion,
+			Handler: func(sess ssh.Session) {
+				uuidSession := uuid.New()
+
+				tr.TraceEvent(tracer.Event{
+					Msg:         "New SSH Session",
+					Protocol:    tracer.SSH.String(),
+					RemoteAddr:  sess.RemoteAddr().String(),
+					Status:      tracer.Start.String(),
+					ID:          uuidSession.String(),
+					Environ:     strings.Join(sess.Environ(), ","),
+					User:        sess.User(),
+					Description: beelzebubServiceConfiguration.Description,
+					Command:     sess.RawCommand(),
+				})
+
+				term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
+				var histories []plugins.History
+				for {
+					commandInput, err := term.ReadLine()
+					if err != nil {
+						break
+					}
+					tr.TraceEvent(tracer.Event{
+						Msg:         "New SSH Terminal Session",
+						RemoteAddr:  sess.RemoteAddr().String(),
+						Status:      tracer.Interaction.String(),
+						Command:     commandInput,
+						ID:          uuidSession.String(),
+						Protocol:    tracer.SSH.String(),
+						Description: beelzebubServiceConfiguration.Description,
+					})
+					if commandInput == "exit" {
+						break
+					}
+					for _, command := range beelzebubServiceConfiguration.Commands {
+						matched, err := regexp.MatchString(command.Regex, commandInput)
+						if err != nil {
+							log.Errorf("Error regex: %s, %s", command.Regex, err.Error())
+							continue
+						}
+
+						if matched {
+							commandOutput := command.Handler
+
+							if command.Plugin == plugins.ChatGPTPluginName {
+								openAIGPTVirtualTerminal := plugins.OpenAIGPTVirtualTerminal{Histories: histories, OpenAPIChatGPTSecretKey: beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey}
+								openAIGPTVirtualTerminal.InjectDependency()
+
+								if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
+									log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
+									commandOutput = "command not found"
+								}
+							}
+
+							histories = append(histories, plugins.History{Input: commandInput, Output: commandOutput})
+
+							term.Write(append([]byte(commandOutput), '\n'))
+							break
+						}
+					}
+				}
+				tr.TraceEvent(tracer.Event{
+					Msg:    "End SSH Session",
+					Status: tracer.End.String(),
+					ID:     uuidSession.String(),
+				})
+			},
+			PasswordHandler: func(ctx ssh.Context, password string) bool {
+				tr.TraceEvent(tracer.Event{
+					Msg:         "New SSH attempt",
+					Protocol:    tracer.SSH.String(),
+					Status:      tracer.Stateless.String(),
+					User:        ctx.User(),
+					Password:    password,
+					Client:      ctx.ClientVersion(),
+					RemoteAddr:  ctx.RemoteAddr().String(),
+					ID:          uuid.New().String(),
+					Description: beelzebubServiceConfiguration.Description,
+				})
+				matched, err := regexp.MatchString(beelzebubServiceConfiguration.PasswordRegex, password)
+				if err != nil {
+					log.Errorf("Error regex: %s, %s", beelzebubServiceConfiguration.PasswordRegex, err.Error())
+					return false
+				}
+				return matched
+			},
+		}
+		err := server.ListenAndServe()
+		if err != nil {
+			log.Errorf("Error during init SSH Protocol: %s", err.Error())
+		}
+	}()
+
+	log.WithFields(log.Fields{
+		"port":     beelzebubServiceConfiguration.Address,
+		"commands": len(beelzebubServiceConfiguration.Commands),
+	}).Infof("Init service %s", beelzebubServiceConfiguration.Protocol)
+	return nil
+}
+
+func buildPrompt(user string, serverName string) string {
+	return fmt.Sprintf("%s@%s:~$ ", user, serverName)
+}