main.go

// main.go

package main

import (
	"encoding/hex"
	"fmt"
	"os"
	"strings"

	"github.com/charmbracelet/bubbles/textinput"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

// START_STYLES

var (
	titleStyle        = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("#FAFAFA")).Background(lipgloss.Color("#7D56F4")).Padding(0, 1).MarginBottom(1)
	itemStyle         = lipgloss.NewStyle().PaddingLeft(2)
	selectedItemStyle = lipgloss.NewStyle().PaddingLeft(2).Foreground(lipgloss.Color("205")).Bold(true)

	// INPUT STYLES
	focusedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("205"))
	blurredStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240"))
	noStyle      = lipgloss.NewStyle()
)

// END_STYLES

// START_APPLICATION_STATE
type screen int

const (
	screenMenu screen = iota
	screenRegister
	screenStore
	screenRead
	screenTamper
	screenResult
)

// END_APPLICATION_STATE

// START_BUBBLE_TEA_MODEL
type model struct {
	screen  screen
	cursor  int
	choices []string

	// REGISTRATION INPUT
	inputs     []textinput.Model
	focusIndex int

	// STORE NOTE INPUTS
	storeInputs     []textinput.Model
	storeFocusIndex int

	// READ NOTE INPUTS
	readInputs     []textinput.Model
	readFocusIndex int

	// RESULT MESSAGE
	resultMsg string
}

func initialModel() model {
	// REGISTRATION FORM
	// USERNAME
	inputs := make([]textinput.Model, 3)
	inputs[0] = textinput.New()
	inputs[0].Placeholder = "Please enter username"
	inputs[0].Focus()
	inputs[0].Prompt = "Username: "
	inputs[0].CharLimit = 20

	// PASSWORD
	inputs[1] = textinput.New()
	inputs[1].Placeholder = "Please enter password"
	inputs[1].EchoMode = textinput.EchoPassword
	inputs[1].Prompt = "Password: "
	inputs[1].CharLimit = 20

	// HASH MODE
	inputs[2] = textinput.New()
	inputs[2].Placeholder = "SHA256"
	inputs[2].Prompt = "Hash Mode (SHA256): "
	inputs[2].CharLimit = 6

	// STORE NOTE FORM
	// USERNAME
	storeInputs := make([]textinput.Model, 3)
	storeInputs[0] = textinput.New()
	storeInputs[0].Placeholder = "Please enter username"
	// storeInputs[0].Focus()
	storeInputs[0].Prompt = "Username: "
	// storeInputs[0].CharLimit = 20

	// PASSWORD
	storeInputs[1] = textinput.New()
	storeInputs[1].Placeholder = "Please enter password"
	storeInputs[1].EchoMode = textinput.EchoPassword
	storeInputs[1].Prompt = "Password: "
	// storeInputs[1].CharLimit = 20

	// SECRET NOTE
	storeInputs[2] = textinput.New()
	storeInputs[2].Placeholder = "Type your secret note here"
	storeInputs[2].Prompt = "Note: "
	// storeInputs[2].CharLimit = 6

	// READ NOTE FORM
	readInputs := make([]textinput.Model, 2)
	readInputs[0] = textinput.New()
	readInputs[0].Placeholder = "Please enter username"
	readInputs[0].Prompt = "Username: "

	readInputs[1] = textinput.New()
	readInputs[1].Placeholder = "Please enter password"
	readInputs[1].EchoMode = textinput.EchoPassword
	readInputs[1].Prompt = "Password: "

	return model{
		screen:  screenMenu,
		choices: []string{"Register User", "Store Note", "Read Note", "Tamper Data", "Delete User", "Exit"},
		cursor:  0, inputs: inputs, focusIndex: 0, storeInputs: storeInputs, storeFocusIndex: 0, readInputs: readInputs,
	}
}

func (m model) Init() tea.Cmd {
	return textinput.Blink
}

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

		switch msg.String() {

		// GLOBAL QUIT KEYS
		case "ctrl+c", "q":
			if m.screen == screenMenu {
				return m, tea.Quit
			}

		case "esc":
			m.screen = screenMenu
			return m, nil

			// MENU NAVIGATION
		case "enter", " ":
			switch m.screen {

			case screenMenu:
				switch m.cursor {
				case 0:
					m.screen = screenRegister
					m.inputs[0].Focus()
				case 1:
					m.screen = screenStore
					m.storeInputs[0].Focus()
				case 2:
					m.screen = screenRead
					m.readInputs[0].Focus()
				case 3:
					vaultData, err := loadVault()

					if err != nil || vaultData.Ciphertext == "" {
						m.resultMsg = "Error: Vault is EMPTY!. Store a note first!"
					} else {
						tamperedText := []byte(vaultData.Ciphertext)

						if tamperedText[0] == 'A' {
							tamperedText[0] = 'B'
						} else {
							tamperedText[0] = 'A'
						}

						vaultData.Ciphertext = string(tamperedText)

						if err := saveVault(vaultData); err != nil {
							m.resultMsg = "Error saving tampered vault: " + err.Error()
						} else {
							m.resultMsg = "HEHE! DATA TAMPERED SUCCESSFULLY ;)\n\nA byte in the ciphertext was intentionally changed.\nTry using 'Read Note' to see if the Integrity Check catches it!\n\n(Press ENTER to Return to menu)"
						}
					}
					m.screen = screenResult
				case 4:
					err := os.Remove("vault.json")

					if err != nil {
						if os.IsNotExist(err) {
							m.resultMsg = "No User found to delete. The Vault is already empty."
						} else {
							m.resultMsg = "Error deleting user: " + err.Error()
						}
					} else {
						m.resultMsg = "User and all vault data deleted successfully!\n\n(Press ENTER to return to menu)"
					}

					m.screen = screenResult
				case 5:
					return m, tea.Quit
				}
			case screenResult:
				m.screen = screenMenu

			case screenRegister:
				// FORM SUBMISSION
				if m.focusIndex == len(m.inputs)-1 {
					// REGISTRATION TRIGGER LOGIC
					username := m.inputs[0].Value()
					// password := m.inputs[1].Value()
					hashMode := m.inputs[2].Value()

					if hashMode == "" {
						hashMode = "SHA256"
					}

					salt, err := generateRandomBytes(16)

					if err != nil {
						m.resultMsg = "Error generating salt: " + err.Error()
					} else {
						vaultData := VaultData{
							Username: username,
							Salt:     salt, HashMode: hashMode,
						}
						if err := saveVault(vaultData); err != nil {
							m.resultMsg = "Error saving vault: " + err.Error()
						} else {
							m.resultMsg = fmt.Sprintf("User Registered Successfully!\n\nGenerated Salt: %s\n\n(Press ENTER to return to menu)", salt)
						}
					}

					m.screen = screenResult

					m.inputs[0].SetValue("")
					m.inputs[1].SetValue("")
					m.inputs[2].SetValue("")
					m.focusIndex = 0
				} else {
					m.focusIndex++
				}
			case screenStore:
				if m.storeFocusIndex == len(m.storeInputs)-1 {
					// STORE NOTE LOGIC
					username := m.storeInputs[0].Value()
					password := m.storeInputs[1].Value()
					note := m.storeInputs[2].Value()

					vaultData, err := loadVault()

					if err != nil || vaultData.Username != username {
						m.resultMsg = "Error: User not found or vault is empty."
					} else {
						// DERIVE KEY
						key, err := deriveKey(password, vaultData.Salt, vaultData.HashMode)
						if err != nil {
							m.resultMsg = "Error deriving key: " + err.Error()
						} else {
							// ENCRYPT NOTE
							iv, ciphertext, err := encryptNote(key, note)

							if err != nil {
								m.resultMsg = "Error encrypting note: " + err.Error()
							} else {
								// COMPUTE INTEGRITY HASH
								integrityHash := computeIntegrityHash(iv, ciphertext, vaultData.HashMode)

								// SAVE TO VAULT
								vaultData.IV = iv
								vaultData.Ciphertext = ciphertext
								vaultData.IntegrityHash = integrityHash
								saveVault(vaultData)
								// BUILD RESULT MESSAGE
								keyHex := hex.EncodeToString(key)[:32]
								m.resultMsg = fmt.Sprintf("Encrypted and stored.\n\n[DEBUG INFO]\nDerived Key (first 16 bytes): %s\n(What is Derived Key? A fixed-length cryptographic key created from the human password and salt.)\n\nIV: %s\n(What is IV? A random value used to randomize encryption outputs.)\n\nCiphertext: %s\n\n(Press ENTER to return to menu)", keyHex, iv, ciphertext)
							}
						}
					}
					m.screen = screenResult
					m.storeInputs[0].SetValue("")
					m.storeInputs[1].SetValue("")
					m.storeInputs[2].SetValue("")
					m.storeFocusIndex = 0
				} else {
					m.storeFocusIndex++
				}

			case screenRead:
				if m.readFocusIndex == len(m.readInputs)-1 {
					// READ NOTE LOGIC
					username := m.readInputs[0].Value()
					password := m.readInputs[1].Value()

					vaultData, err := loadVault()

					if err != nil || vaultData.Username != username {
						m.resultMsg = "Error: User not found or vault if empty."
					} else if vaultData.Ciphertext == "" {
						m.resultMsg = "Error: No note stored in vault yet."
					} else {
						// INTEGRITY CHECK
						calculatedHash := computeIntegrityHash(vaultData.IV, vaultData.Ciphertext, vaultData.HashMode)

						if calculatedHash != vaultData.IntegrityHash {
							// HASH MISMATCH, TAMPERING DETECTED
							m.resultMsg = "INTEGRITY CHECK FAILED!\nStored Hash: " + vaultData.IntegrityHash + "\nCalculated: " + calculatedHash + "\n\nThe vault data has been tampered with or corrupted!"
						} else {
							// DERIVE KEY
							key, err := deriveKey(password, vaultData.Salt, vaultData.HashMode)

							if err != nil {
								m.resultMsg = "Error deriving key: " + err.Error()
							} else {

								// DECRYPT
								plaintext, err := decryptNote(key, vaultData.IV, vaultData.Ciphertext)

								if err != nil {
									m.resultMsg = "Decryption failed! Incorrect password or corrupted data.\nError: " + err.Error()
								} else {
									m.resultMsg = fmt.Sprintf("Integrity Check: PASSED!\n\nSECRET NOTE: \n%s\n\n(Press ENTER to return to menu)", plaintext)
								}
							}
						}
					}

					m.screen = screenResult
					m.readInputs[0].SetValue("")
					m.readInputs[1].SetValue("")
					m.readFocusIndex = 0
				} else {
					m.readFocusIndex++
				}
			}

		case "up", "k":
			if m.screen == screenMenu && m.cursor > 0 {
				m.cursor--
			} else if m.screen == screenRegister && m.focusIndex > 0 {
				m.focusIndex--
			} else if m.screen == screenStore && m.storeFocusIndex > 0 {
				m.storeFocusIndex--
			} else if m.screen == screenRead && m.readFocusIndex > 0 {
				m.readFocusIndex--
			}

		case "down", "j":
			if m.screen == screenMenu && m.cursor < len(m.choices)-1 {
				m.cursor++
			} else if m.screen == screenRegister && m.focusIndex < len(m.inputs)-1 {
				m.focusIndex++
			} else if m.screen == screenStore && m.storeFocusIndex < len(m.storeInputs)-1 {
				m.storeFocusIndex++
			} else if m.screen == screenRead && m.readFocusIndex < len(m.readInputs)-1 {
				m.readFocusIndex++
			}
		}
	}

	cmd := m.updateInputs(msg)
	return m, cmd
}

// HELPER: UPDATE INPUT FIELDS
func (m *model) updateInputs(msg tea.Msg) tea.Cmd {
	var cmds []tea.Cmd

	switch m.screen {
	case screenRegister:
		cmds := make([]tea.Cmd, len(m.inputs))

		for i := range m.inputs {
			m.inputs[i].Cursor.Style = noStyle

			if i == m.focusIndex {
				m.inputs[i].Cursor.Style = focusedStyle
				m.inputs[i].TextStyle = focusedStyle
				m.inputs[i].PromptStyle = focusedStyle
				m.inputs[i].Focus()
			} else {
				m.inputs[i].Blur()
				m.inputs[i].TextStyle = noStyle
				m.inputs[i].PromptStyle = noStyle
			}

			m.inputs[i], cmds[i] = m.inputs[i].Update(msg)
		}
	case screenStore:
		cmds = make([]tea.Cmd, len(m.storeInputs))

		for i := range m.storeInputs {
			m.storeInputs[i].Cursor.Style = noStyle

			if i == m.storeFocusIndex {
				m.storeInputs[i].Cursor.Style = focusedStyle
				m.storeInputs[i].TextStyle = focusedStyle
				m.storeInputs[i].PromptStyle = focusedStyle
				m.storeInputs[i].Focus()
			} else {
				m.storeInputs[i].Blur()
				m.storeInputs[i].TextStyle = noStyle
				m.storeInputs[i].PromptStyle = noStyle
			}

			m.storeInputs[i], cmds[i] = m.storeInputs[i].Update(msg)
		}
	case screenRead:
		cmds = make([]tea.Cmd, len(m.readInputs))

		for i := range m.readInputs {
			m.readInputs[i].Cursor.Style = noStyle

			if i == m.readFocusIndex {
				m.readInputs[i].Cursor.Style = focusedStyle
				m.readInputs[i].TextStyle = focusedStyle
				m.readInputs[i].PromptStyle = focusedStyle
				m.readInputs[i].Focus()
			} else {
				m.readInputs[i].Blur()
				m.readInputs[i].TextStyle = noStyle
				m.readInputs[i].PromptStyle = noStyle
			}

			m.readInputs[i], cmds[i] = m.readInputs[i].Update(msg)
		}
	}

	return tea.Batch(cmds...)
}

func (m model) View() string {
	var b strings.Builder

	b.WriteString("\n")
	b.WriteString(titleStyle.Render("CS6413 Assignment#1 - Secure Drop Vault") + "\n")

	// RENDER MENU
	switch m.screen {

	case screenMenu:
		for i, choice := range m.choices {
			cursor := " "
			if m.cursor == i {
				cursor = ">"
				b.WriteString(selectedItemStyle.Render(fmt.Sprintf("%s %s", cursor, choice)) + "\n")
			} else {
				b.WriteString(itemStyle.Render(fmt.Sprintf("%s %s", cursor, choice)) + "\n")
			}
		}

		b.WriteString("\n(Press q to quit)\n")

	case screenRegister:
		b.WriteString(lipgloss.NewStyle().Bold(true).MarginBottom(1).Render("USER REGISTRATION") + "\n")

		for i := range m.inputs {
			b.WriteString(m.inputs[i].View() + "\n")
		}

		b.WriteString("\n(Press ESC to Cancel, ENTER to Submit)\n")

	case screenStore:
		b.WriteString(lipgloss.NewStyle().Bold(true).MarginBottom(1).Render("STORE ENCRYPT NOTE") + "\n")

		for i := range m.storeInputs {
			b.WriteString(m.storeInputs[i].View() + "\n")
		}

		b.WriteString("\n(Press ESC to Cancel, ENTER to Submit)\n")

	case screenRead:
		b.WriteString(lipgloss.NewStyle().Bold(true).MarginBottom(1).Render("READ SECURE NOTE") + "\n")

		for i := range m.readInputs {
			b.WriteString((m.readInputs[i].View() + "\n"))
		}

		b.WriteString("\n(Press ESC to Cancel, ENTER to Submit)\n")

	case screenResult:
		b.WriteString(lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).Padding(1, 2).Foreground(lipgloss.Color("205")).Render(m.resultMsg) + "\n")
	}

	return b.String()
}

// END_BUBBLE_TEA_MODEL

func main() {
	p := tea.NewProgram(initialModel())
	if _, err := p.Run(); err != nil {
		fmt.Printf("Error starting program: %v", err)
		os.Exit(1)
	}
}