protomolecule/tui.go

package main

/*
 */

import (
	"flag"
	"fmt"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/mattn/go-runewidth"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	//"io"
	"os"
	"protomolecule/src/dust"
	"protomolecule/src/eros"
	"protomolecule/src/scanStuff"
	projVars "protomolecule/src/vars"
	"strings"
	"time"
)

var _ScanMgr *scanStuff.Meta

const (
	headerHeight = 3
	footerHeight = 3
)

type model struct {

	// https://github.com/charmbracelet/bubbles/blob/master/viewport/viewport.go
	viewport viewport.Model
	content  string

	// items on the device list
	choices []string
	// which device item our cursor is pointing at
	cursor int
	// which device items are selected
	selected map[int]struct{}
	// have we determined the window dimensions yet
	ready bool
	// it spins
	spinner spinner.Model
}

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{"Address" + "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,
}

//var reader *io.PipeReader
//var writer *io.PipeWriter

func init() {
	//reader, writer = io.Pipe()

	// 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: writer,
	//}

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

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

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

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

}

// this tomfoolery was taken from the example for viewports
// i think the reasoning is aligned with the dynamic initialization that comes with this type of model
//
// tldr; initialization happens repeatedly in the update function
func (m model) Init() tea.Cmd {
	return nil
}

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

// Update is processing "msg"s (messages) which appear to be bubbletea's utilization of channels for IPC
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

	var (
		//cmd  tea.Cmd
		cmds []tea.Cmd
	)

	cmds = append(cmds, viewport.Sync(m.viewport))

	switch msg := msg.(type) {

	case tea.WindowSizeMsg:
		verticalMargins := headerHeight + footerHeight

		// getting the size of the window happens asynchronously so
		// we wait to receive them (ready) to determine the viewport dimensions
		if !m.ready {
			m.viewport = viewport.Model{
				Width:  msg.Width,
				Height: msg.Height - verticalMargins,
			}

			m.viewport.YPosition = headerHeight
			m.viewport.HighPerformanceRendering = true

			// so it looks like this is fetching data from our io.Pipe (io.reader) for every update
			// i defined said pipe on line 65
			m.viewport.SetContent(m.content)

			m.ready = true
		} else {
			// now that we've figured out window dimensions we define viewport dimensions
			m.viewport.Width = msg.Width
			m.viewport.Height = msg.Height - verticalMargins
		}

		cmds = append(cmds, viewport.Sync(m.viewport))

	// 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{}{}
			}
		}

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

	// **************************************************
	// actual protomolecule IPC for bubbletea starts here
	// **************************************************

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

				cmds = append(cmds, listenForScanResults)

				return m, tea.Batch(cmds...)
			}
		}

		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, tea.Batch(cmds...)
}

func (m model) View() string {
	//var head []string

	if !m.ready {
		return "Loading ProtoMolecule..."
	}

	// load our banner
	//	head = dust.GetBanner()

	header := lipgloss.NewStyle().
		Bold(true).
		Foreground(lipgloss.Color("#5991D4")).
		Align(lipgloss.Center).
		PaddingTop(0).
		PaddingLeft(0).
		Width(m.viewport.Width).
		Render("ProtoMolecule")
		//		Render(fmt.Sprintf("\n%s\n%s\n%s", head[0], head[1], head[2]))

	footerTop := "╭──────╮"
	footerMid := fmt.Sprintf("┤ %3.f%% │", m.viewport.ScrollPercent()*100)
	footerBot := "╰──────╯"
	gapSize := m.viewport.Width - runewidth.StringWidth(footerMid)
	footerTop = strings.Repeat(" ", gapSize) + footerTop
	footerMid = strings.Repeat("─", gapSize) + footerMid
	footerBot = strings.Repeat(" ", gapSize) + footerBot
	footer := fmt.Sprintf("%s\n%s\n%s", footerTop, footerMid, footerBot)

	/* 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, " ") {
			m.content += 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 = lipgloss.
				NewStyle().
				Bold(true).
				Render("x")
		}

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

	// Send the UI for rendering
	return fmt.Sprintf("%s\n%s\n%s", header, m.viewport.View(), footer)
}

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

func main() {
	//content := "Starting..."

	p := tea.NewProgram(initialModel)
	p.EnterAltScreen()
	defer p.ExitAltScreen()

	log.Debug().Msg("Starting scan")

	var scanID int
	var scan *scanStuff.Scan

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

	go scan.Start()

	if err := p.Start(); err != nil {
		log.Fatal().Err(err).Msg("FATAL")
	}
}