oubliette/internal/shell/bash/bash.go

package bash

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"code.t-juice.club/torjus/oubliette/internal/shell"
)

const sessionTimeout = 5 * time.Minute

// BashShell emulates a basic bash-like shell.
type BashShell struct{}

// NewBashShell returns a new BashShell instance.
func NewBashShell() *BashShell {
	return &BashShell{}
}

func (b *BashShell) Name() string        { return "bash" }
func (b *BashShell) Description() string { return "Basic bash-like shell emulator" }

func (b *BashShell) Handle(ctx context.Context, sess *shell.SessionContext, rw io.ReadWriteCloser) error {
	ctx, cancel := context.WithTimeout(ctx, sessionTimeout)
	defer cancel()

	username := sess.Username
	if sess.CommonConfig.FakeUser != "" {
		username = sess.CommonConfig.FakeUser
	}
	hostname := sess.CommonConfig.Hostname

	fs := newFilesystem(hostname)
	state := &shellState{
		cwd:      "/root",
		username: username,
		hostname: hostname,
		fs:       fs,
	}

	// Send banner.
	if sess.CommonConfig.Banner != "" {
		fmt.Fprint(rw, sess.CommonConfig.Banner)
	}
	fmt.Fprintf(rw, "Last login: %s from 10.0.0.1\r\n",
		time.Now().Add(-2*time.Hour).Format("Mon Jan 2 15:04:05 2006"))

	for {
		prompt := formatPrompt(state)
		if _, err := fmt.Fprint(rw, prompt); err != nil {
			return nil
		}

		line, err := shell.ReadLine(ctx, rw)
		if errors.Is(err, io.EOF) {
			fmt.Fprint(rw, "logout\r\n")
			return nil
		}
		if err != nil {
			return nil
		}

		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}

		result := dispatch(state, trimmed)

		var output string
		if result.output != "" {
			output = result.output
			// Convert newlines to \r\n for terminal display.
			output = strings.ReplaceAll(output, "\r\n", "\n")
			output = strings.ReplaceAll(output, "\n", "\r\n")
			fmt.Fprintf(rw, "%s\r\n", output)
		}

		// Log command and output to store.
		if sess.Store != nil {
			if err := sess.Store.AppendSessionLog(ctx, sess.SessionID, trimmed, output); err != nil {
				return fmt.Errorf("append session log: %w", err)
			}
		}
		if sess.OnCommand != nil {
			sess.OnCommand("bash")
		}

		if result.exit {
			return nil
		}
	}
}

func formatPrompt(state *shellState) string {
	cwd := state.cwd
	if cwd == "/root" {
		cwd = "~"
	} else if strings.HasPrefix(cwd, "/root/") {
		cwd = "~" + cwd[5:]
	}
	return fmt.Sprintf("%s@%s:%s# ", state.username, state.hostname, cwd)
}