oubliette/internal/shell/bash/bash.go

package bash

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

	"git.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 := readLine(ctx, rw)
		if 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 {
			sess.Store.AppendSessionLog(ctx, sess.SessionID, trimmed, output)
		}

		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)
}

// readLine reads a line of input byte-by-byte, handling backspace, Ctrl+C, and Ctrl+D.
func readLine(ctx context.Context, rw io.ReadWriter) (string, error) {
	var buf []byte
	b := make([]byte, 1)

	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
		}

		n, err := rw.Read(b)
		if err != nil {
			return "", err
		}
		if n == 0 {
			continue
		}

		ch := b[0]
		switch {
		case ch == '\r' || ch == '\n':
			fmt.Fprint(rw, "\r\n")
			return string(buf), nil

		case ch == 4: // Ctrl+D
			if len(buf) == 0 {
				return "", io.EOF
			}

		case ch == 3: // Ctrl+C
			fmt.Fprint(rw, "^C\r\n")
			return "", nil

		case ch == 127 || ch == 8: // DEL or Backspace
			if len(buf) > 0 {
				buf = buf[:len(buf)-1]
				fmt.Fprint(rw, "\b \b")
			}

		case ch == 27: // ESC - start of escape sequence
			// Read and discard the rest of the escape sequence.
			// Most are 3 bytes: ESC [ X (arrow keys, etc.)
			next := make([]byte, 1)
			rw.Read(next)
			if next[0] == '[' {
				rw.Read(next) // read the final byte
			}

		case ch >= 32 && ch < 127: // printable ASCII
			buf = append(buf, ch)
			rw.Write([]byte{ch})
		}
	}
}