protomolecule/tui.go

package main

import (
	"flag"
	"fmt"
	"github.com/charmbracelet/bubbles/spinner"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	term "github.com/muesli/termenv"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"os"
	"protomolecule/src/dust"
	"protomolecule/src/eros"
	"protomolecule/src/scanStuff"
	projVars "protomolecule/src/vars"
	"strings"
	"time"
)

var _ScanMgr *scanStuff.Meta

type model struct {
	choices  []string         // items on the device list
	cursor   int              // which device item our cursor is pointing at
	selected map[int]struct{} // which device items are selected
	spinner  spinner.Model
}

func init() {

}

func getSpinner() spinner.Model {
	s := spinner.NewModel()
	s.Spinner = spinner.Dot
	s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("205"))

	return s
}

var initialModel = model{
	choices: []string{strings.Repeat(" ", 6) + "Address" + strings.Repeat(" ", 17) + "RSSI"},
	// A map which indicates which choices are selected. We're using
	// the  map like a mathematical set. The keys refer to the indexes
	// of the `choices` slice, above.
	selected: make(map[int]struct{}),
	spinner:  getSpinner(),
	cursor:   1,
}

func init() {
	term.ClearScreen()

	// print banner for style points
	// dust.Splash()

	flag.Parse()

	if *projVars.AFlag {
		projVars.AttackMode = true
	}

	if *projVars.DFlag {
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	}

	if *projVars.TFlag {
		projVars.TrackingMode = true
	}

	_ScanMgr = &scanStuff.Meta{
		Count: 0,
		Scans: make(map[int]*scanStuff.Scan),
	}

	// TODO: make this a commandline argument
	// assure the log directory exists
	var logDir string = "./.logs/"

	err := os.MkdirAll(logDir, 0755)
	if err != nil {
		panic(err.Error())
	}

	// define log file itself using the current date and time
	Now := time.Now()
	date := Now.Format(time.RFC3339)
	logFileName := date + ".log"
	lf, err := os.OpenFile(logDir+logFileName, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(err.Error())
	}

	// define pretty printer
	consoleWriter := zerolog.ConsoleWriter{
		Out: os.Stderr,
	}

	// initialize simultaneous pretty printing and json logging
	multi := zerolog.MultiLevelWriter(consoleWriter, lf)
	log.Logger = zerolog.New(multi).With().Timestamp().Logger()

	// suppress debug messages unless -d is called
	zerolog.SetGlobalLevel(zerolog.Disabled)

	log.Debug().Msg("Logging initialized")

	log.Debug().Msg("Initializing database engine")
	// initialize database engine
	eros.Awaken()
}

func (m model) Init() tea.Cmd {

	return tea.Batch(
		func() tea.Msg {
			var scanID int
			var scan *scanStuff.Scan

			scanID = _ScanMgr.NewScan()
			scan = _ScanMgr.Scans[scanID]

			//time.Sleep(30 * time.Millisecond)
			dust.Must("Scan", scan.Start())

			return nil
		},
		listenForScanResults,
		spinner.Tick,
	)
}

func getFormattedRow(dev projVars.DiscoveredDevice) string {
	return fmt.Sprintf("%s%s%s",
		dev.ScanResult.Address.String(),
		strings.Repeat(" ", 7),
		func(rssi int16) string {
			if rssi < -80 {
				return lipgloss.NewStyle().Foreground(lipgloss.Color("#FF0000")).Render(fmt.Sprintf("%d", rssi))
			}
			if rssi < -50 {
				return lipgloss.NewStyle().Foreground(lipgloss.Color("#FFFF00")).Render(fmt.Sprintf("%d", rssi))
			} else {
				return lipgloss.NewStyle().Foreground(lipgloss.Color("#008000")).Render(fmt.Sprintf("%d", rssi))
			}
		}(dev.ScanResult.RSSI))
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {

	case projVars.DiscoveredDevice:
		for k, v := range m.choices {
			if strings.HasPrefix(v, msg.ScanResult.Address.String()) {
				m.choices[k] = getFormattedRow(msg)

				return m, listenForScanResults
			}
		}

		/* m.choices = append(m.choices[func(lLength int) int {
			if lLength >= 10 {
				return 1
			} else {
				return 0
			}
		}(len(m.choices)):],
			fmt.Sprintf("address: %s rssi: %d", msg.ScanResult.Address.String(), msg.ScanResult.RSSI)) */

		m.choices = append(m.choices, getFormattedRow(msg))

		return m, listenForScanResults

	case spinner.TickMsg:
		var cmd tea.Cmd
		m.spinner, cmd = m.spinner.Update(msg)
		return m, cmd

	// Is it a key press?
	case tea.KeyMsg:

		// Cool, what was the actual key pressed?
		switch msg.String() {

		// These keys should exit the program.
		case "ctrl+c", "q":
			return m, tea.Quit

		// The "up" and "k" keys move the cursor up
		case "up", "k":
			if m.cursor > 1 {
				m.cursor--
			}

		// The "down" and "j" keys move the cursor down
		case "down", "j":
			if m.cursor < len(m.choices)-1 {
				m.cursor++
			}

		// The "enter" key and the spacebar (a literal space) toggle
		// the selected state for the item that the cursor is pointing at.
		case "enter", " ":
			_, ok := m.selected[m.cursor]
			if ok {
				delete(m.selected, m.cursor)
			} else {
				m.selected[m.cursor] = struct{}{}
			}
		}
	}

	// Return the updated model to the Bubble Tea runtime for processing.
	// Note that we're not returning a command.
	return m, nil
}

func (m model) View() string {
	// The header
	s := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("#FAFAFA")).Background(lipgloss.Color("#7D56F4")).PaddingTop(0).PaddingLeft(0).Width(78).Render("protomolecule"+strings.Repeat(" ", 78-24)+"[q] - exit") + "\n\n"

	/* Iterate over our choices
	NOTE: change to only print n rows (minus header, footer and any padding rows) to ensure
	that the number of printed rows is not overlapping the number of blank rows available in
	the terminal; this will cause banding issues as most rows are updated on the screen without
	being completely re-printed (\r vs \n) Scrolling can be added to the Update().
	https://github.com/muesli/termenv
	https://stackoverflow.com/questions/16569433/get-terminal-size-in-go
	https://stackoverflow.com/questions/24562942/golang-how-do-i-determine-the-number-of-lines-in-a-file-efficiently/24563853
	*/
	for i, choice := range m.choices {
		if strings.HasPrefix(choice, " ") {
			s += fmt.Sprintf("%s\n", choice)
			continue
		}

		// Is the cursor pointing at this choice?
		cursor := " " // no cursor
		if m.cursor == i {
			//cursor = ">" // cursor!
			cursor = lipgloss.NewStyle().Blink(true).Render(">")
		}

		// Is this choice selected?
		checked := " " // not selected
		if _, ok := m.selected[i]; ok {
			//checked = "x" // selected!
			checked = lipgloss.NewStyle().Bold(true).Render("x")
		}

		// Render the row
		s += fmt.Sprintf("%s [%s] %s\n", cursor, checked, choice)
	}

	// The footer
	s += m.spinner.View()

	// Send the UI for rendering
	return s
}

func listenForScanResults() tea.Msg {
	for {
		select {
		case dev, ok := <-projVars.DiscoveredDeviceChan:
			if ok {
				//time.Sleep( 1 * time.Second)
				return dev
			} else {
				time.Sleep(1 * time.Second)
			}
		}
	}
}

func main() {
	defer term.ClearScreen()
	p := tea.NewProgram(initialModel)

	if err := p.Start(); err != nil {
		fmt.Printf("Alas, there's been an error: %v", err)
		os.Exit(1)
	}
}