From 2e30805280025d6c1b5a6fc34d6cd1ad12a90fa7 Mon Sep 17 00:00:00 2001 From: Debosmit Ray Date: Mon, 7 Oct 2024 14:43:39 +0000 Subject: [PATCH] add support to instrument all shells at install time --- cmd/cmd.go | 187 +++++++++++++++++++++--------------- config/config.go | 1 - config/os.go | 50 +++++----- daemon/daemon.go | 17 ++-- daemon/services/lda.plist | 2 - daemon/services/lda.service | 1 - database/migrations.go | 25 ++++- user/user.go | 108 ++++++++++++++------- 8 files changed, 242 insertions(+), 149 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 2385c3d..eb137d6 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -14,6 +14,7 @@ import ( "lda/util" "net/http" "os" + "strings" "github.com/pkg/errors" "github.com/spf13/cobra" @@ -31,12 +32,7 @@ var ( Run: lda, } - installCmd = &cobra.Command{ - Use: "install", - Short: "Install daemon runner", - Long: `Install daemon runner for LDA Project.`, - RunE: install, - } + installCmd = newInstallCmd() uninstallCmd = &cobra.Command{ Use: "uninstall", @@ -81,6 +77,23 @@ var ( } ) +var installFlags struct { + shells []string +} + +func newInstallCmd() *cobra.Command { + installCmd := &cobra.Command{ + Use: "install", + Short: "Install daemon runner", + Long: `Install daemon runner for LDA Project.`, + RunE: install, + } + + installCmd.Flags().StringSliceVarP(&installFlags.shells, "shell", "s", []string{}, fmt.Sprintf("Shells to instrument %+v; --shell=all for all shells", config.SupportedShells)) + + return installCmd +} + const ( // days are amount of days that old data will be retained days = 5 @@ -121,7 +134,6 @@ func includeShowFlagsForInstall(cmd *cobra.Command) { } func setupConfig() { - // setting up the system configuration config.SetupSysConfig() @@ -203,12 +215,12 @@ func reload(_ *cobra.Command, _ []string) error { user.ConfigureUserSystemInfo(user.Conf) daemonConf := &daemon.Config{ - ExePath: user.Conf.ExePath, - HomeDir: user.Conf.HomeDir, - IsRoot: user.Conf.IsRoot, - Os: config.OSType(user.Conf.Os), - SudoExecUser: user.Conf.User, - ShellLocation: user.Conf.ShellLocation, + ExePath: user.Conf.ExePath, + HomeDir: user.Conf.HomeDir, + IsRoot: user.Conf.IsRoot, + Os: config.OSType(user.Conf.Os), + SudoExecUser: user.Conf.User, + ShellTypeToLocation: user.Conf.ShellTypeToLocation, } dmn := daemon.NewDaemon(daemonConf, logging.Log) @@ -226,12 +238,12 @@ func start(_ *cobra.Command, _ []string) error { user.ConfigureUserSystemInfo(user.Conf) daemonConf := &daemon.Config{ - ExePath: user.Conf.ExePath, - HomeDir: user.Conf.HomeDir, - IsRoot: user.Conf.IsRoot, - Os: config.OSType(user.Conf.Os), - SudoExecUser: user.Conf.User, - ShellLocation: user.Conf.ShellLocation, + ExePath: user.Conf.ExePath, + HomeDir: user.Conf.HomeDir, + IsRoot: user.Conf.IsRoot, + Os: config.OSType(user.Conf.Os), + SudoExecUser: user.Conf.User, + ShellTypeToLocation: user.Conf.ShellTypeToLocation, } dmn := daemon.NewDaemon(daemonConf, logging.Log) @@ -249,12 +261,12 @@ func stop(_ *cobra.Command, _ []string) error { user.ConfigureUserSystemInfo(user.Conf) daemonConf := &daemon.Config{ - ExePath: user.Conf.ExePath, - HomeDir: user.Conf.HomeDir, - IsRoot: user.Conf.IsRoot, - Os: config.OSType(user.Conf.Os), - SudoExecUser: user.Conf.User, - ShellLocation: user.Conf.ShellLocation, + ExePath: user.Conf.ExePath, + HomeDir: user.Conf.HomeDir, + IsRoot: user.Conf.IsRoot, + Os: config.OSType(user.Conf.Os), + SudoExecUser: user.Conf.User, + ShellTypeToLocation: user.Conf.ShellTypeToLocation, } dmn := daemon.NewDaemon(daemonConf, logging.Log) @@ -268,6 +280,26 @@ func stop(_ *cobra.Command, _ []string) error { } func install(cmd *cobra.Command, _ []string) error { + // validate the stuff + if len(installFlags.shells) > 0 { + for _, shellType := range installFlags.shells { + if strings.EqualFold(shellType, "all") { + installFlags.shells = config.SupportedShells + break + } else if config.GetShellType(shellType) == -1 { + fmt.Fprintf(config.SysConfig.ErrOut, "Unsupported shell: %s\nPlease choose one of: %+v\n", shellType, config.SupportedShells) + os.Exit(1) + } + } + } + + // if shells are provided, lets deal with it + if len(installFlags.shells) > 0 { + user.Conf.ShellTypeToLocation = make(map[config.ShellType]string) + for _, shell := range installFlags.shells { + user.Conf.ShellTypeToLocation[config.GetShellType(shell)] = shell + } + } autoCredentials, err := cmd.Flags().GetBool("auto-credentials") if err != nil { @@ -283,14 +315,14 @@ func install(cmd *cobra.Command, _ []string) error { user.ConfigureUserSystemInfo(user.Conf) daemonConf := &daemon.Config{ - ExePath: user.Conf.ExePath, - HomeDir: user.Conf.HomeDir, - IsRoot: user.Conf.IsRoot, - Os: config.OSType(user.Conf.Os), - SudoExecUser: user.Conf.User, - ShellLocation: user.Conf.ShellLocation, - AutoCredential: autoCredentials, - IsWorkspace: isWorkspace, + ExePath: user.Conf.ExePath, + HomeDir: user.Conf.HomeDir, + IsRoot: user.Conf.IsRoot, + Os: config.OSType(user.Conf.Os), + SudoExecUser: user.Conf.User, + AutoCredential: autoCredentials, + IsWorkspace: isWorkspace, + ShellTypeToLocation: user.Conf.ShellTypeToLocation, } dmn := daemon.NewDaemon(daemonConf, logging.Log) @@ -300,30 +332,31 @@ func install(cmd *cobra.Command, _ []string) error { return errors.Wrap(err, "failed to install LDA daemon configuration file") } - shellConfig := &shell.Config{ - ShellType: config.ShellType(user.Conf.ShellType), - ShellLocation: user.Conf.ShellLocation, - IsRoot: user.Conf.IsRoot, - SudoExecUser: user.Conf.User, - LdaDir: user.Conf.LdaDir, - HomeDir: user.Conf.HomeDir, - } + for shellType, shellLocation := range user.Conf.ShellTypeToLocation { + shellConfig := &shell.Config{ + ShellType: config.ShellType(shellType), + ShellLocation: shellLocation, + IsRoot: user.Conf.IsRoot, + SudoExecUser: user.Conf.User, + LdaDir: user.Conf.LdaDir, + HomeDir: user.Conf.HomeDir, + } - shl, err := shell.NewShell(shellConfig, logging.Log) + shl, err := shell.NewShell(shellConfig, logging.Log) - if err != nil { - logging.Log.Error().Err(err).Msg("Failed to setup shell") - os.Exit(1) - } + if err != nil { + logging.Log.Error().Err(err).Msg("Failed to setup shell") + os.Exit(1) + } - if err := shl.InstallShellConfiguration(); err != nil { - logging.Log.Error().Err(err).Msg("Failed to install shell configuration") - return errors.Wrap(err, "failed to install LDA shell configuration files") - } + if err := shl.InstallShellConfiguration(); err != nil { + logging.Log.Error().Err(err).Msg("Failed to install shell configuration") + return errors.Wrap(err, "failed to install LDA shell configuration files") + } - if err := shl.InjectShellSource(); err != nil { - logging.Log.Error().Err(err).Msg("Failed to inject shell source") - return errors.Wrap(err, "failed to inject LDA shell source") + if err := shl.InjectShellSource(); err != nil { + logging.Log.Error().Err(err).Msgf("Failed to inject shell source (%s); will reattempt at `start` time", shellLocation) + } } fmt.Fprintln(config.SysConfig.Out, "LDA daemon installed successfully.") @@ -335,28 +368,35 @@ func uninstall(_ *cobra.Command, _ []string) error { user.ConfigureUserSystemInfo(user.Conf) daemonConf := &daemon.Config{ - ExePath: user.Conf.ExePath, - HomeDir: user.Conf.HomeDir, - IsRoot: user.Conf.IsRoot, - Os: config.OSType(user.Conf.Os), - SudoExecUser: user.Conf.User, - ShellLocation: user.Conf.ShellLocation, + ExePath: user.Conf.ExePath, + HomeDir: user.Conf.HomeDir, + IsRoot: user.Conf.IsRoot, + Os: config.OSType(user.Conf.Os), + SudoExecUser: user.Conf.User, + ShellTypeToLocation: user.Conf.ShellTypeToLocation, } dmn := daemon.NewDaemon(daemonConf, logging.Log) - shellConfig := &shell.Config{ - ShellType: config.ShellType(user.Conf.ShellType), - ShellLocation: user.Conf.ShellLocation, - IsRoot: user.Conf.IsRoot, - SudoExecUser: user.Conf.User, - LdaDir: user.Conf.LdaDir, - HomeDir: user.Conf.HomeDir, - } - shl, err := shell.NewShell(shellConfig, logging.Log) + for shellType, shellLocation := range user.Conf.ShellTypeToLocation { + shellConfig := &shell.Config{ + ShellType: config.ShellType(shellType), + ShellLocation: shellLocation, + IsRoot: user.Conf.IsRoot, + SudoExecUser: user.Conf.User, + LdaDir: user.Conf.LdaDir, + HomeDir: user.Conf.HomeDir, + } + shl, err := shell.NewShell(shellConfig, logging.Log) - if err != nil { - logging.Log.Error().Err(err).Msg("Failed to setup shell") - os.Exit(1) + if err != nil { + logging.Log.Error().Err(err).Msg("Failed to setup shell") + os.Exit(1) + } + + if err := shl.DeleteShellConfiguration(); err != nil { + logging.Log.Error().Err(err).Msg("Failed to delete shell configuration") + return errors.Wrap(err, "failed to delete LDA shell configuration files") + } } fmt.Fprintln(config.SysConfig.Out, "Uninstalling LDA daemon...") @@ -365,11 +405,6 @@ func uninstall(_ *cobra.Command, _ []string) error { return errors.Wrap(err, "failed to uninstall LDA daemon configuration file") } - if err := shl.DeleteShellConfiguration(); err != nil { - logging.Log.Error().Err(err).Msg("Failed to delete shell configuration") - return errors.Wrap(err, "failed to delete LDA shell configuration files") - } - fmt.Fprintln(config.SysConfig.Out, `Daemon service files and shell configuration deleted successfully, ~/.lda directory still holds database file, and your rc file stills has source script. If you wish to remove those, delete them manually`) diff --git a/config/config.go b/config/config.go index ff77f50..4a16c41 100644 --- a/config/config.go +++ b/config/config.go @@ -78,7 +78,6 @@ func SetupSysConfig() { // SetupConfig initialize the configuration instance func SetupConfig(ldaDir string, user *user.User) { - configPath := filepath.Join(ldaDir, "config.toml") if _, err := os.Stat(configPath); errors.Is(err, os.ErrNotExist) { diff --git a/config/os.go b/config/os.go index faaf1ec..5da6fcf 100644 --- a/config/os.go +++ b/config/os.go @@ -12,6 +12,24 @@ import ( "github.com/manifoldco/promptui" ) +var ( + SupportedShells = []string{"/bin/bash", "/bin/zsh", "/bin/fish"} +) + +func GetShellType(shellLocation string) ShellType { + shellType := path.Base(shellLocation) + switch shellType { + case "bash": + return Bash + case "zsh": + return Zsh + case "fish": + return Fish + default: + return -1 + } +} + // ShellType is the type of the shell that is supported type ShellType int @@ -100,44 +118,30 @@ func GetLdaBinaryPath() (string, error) { } // GetShell sets the current active shell and location -func GetShell() (ShellType, string, error) { - +func GetShell() (map[ShellType]string, error) { shellLocation := os.Getenv("SHELL") - return configureShell(shellLocation) } -func configureShell(shellLocation string) (ShellType, string, error) { - shellType := path.Base(shellLocation) - - var shell ShellType - switch shellType { - case "bash": - shell = Bash - case "zsh": - shell = Zsh - case "fish": - shell = Fish - // TODO: consider supporting "sh" and "ash" as well. - default: +func configureShell(shellLocation string) (map[ShellType]string, error) { + shellTypeToLocation := make(map[ShellType]string) + shell := GetShellType(shellLocation) + if shell < 0 { shellLocation, err := promptForShellType() if err != nil { - return -1, "", err + return shellTypeToLocation, err } return configureShell(shellLocation) } - - return shell, shellLocation, nil + shellTypeToLocation[shell] = shellLocation + return shellTypeToLocation, nil } // promptForShellPath prompts the user to confirm the detected shell path or input a new one. func promptForShellType() (string, error) { - - supportedShells := []string{"/bin/bash", "/bin/zsh", "/bin/fish"} - prompt := promptui.Select{ Label: "We detected an unsupported shell, often this could happen because the script was run as sudo. Currently, we support the following shells. Please select one:", - Items: supportedShells, + Items: SupportedShells, } _, result, err := prompt.Run() diff --git a/daemon/daemon.go b/daemon/daemon.go index 3b01bd5..eedcfed 100644 --- a/daemon/daemon.go +++ b/daemon/daemon.go @@ -37,14 +37,14 @@ var templateFS embed.FS // Config is the configuration for the daemon service type Config struct { - ExePath string - ShellLocation string - HomeDir string - Os config.OSType - IsRoot bool - SudoExecUser *user.User - AutoCredential bool - IsWorkspace bool + ExePath string + HomeDir string + Os config.OSType + IsRoot bool + SudoExecUser *user.User + AutoCredential bool + IsWorkspace bool + ShellTypeToLocation map[config.ShellType]string } // Daemon is the service that configures background service @@ -79,7 +79,6 @@ func (d *Daemon) InstallDaemonConfiguration() error { var content bytes.Buffer var tmpConf = map[string]interface{}{ "BinaryPath": d.config.ExePath, - "Shell": d.config.ShellLocation, "Home": d.config.HomeDir, } diff --git a/daemon/services/lda.plist b/daemon/services/lda.plist index aea3e92..25a59fd 100644 --- a/daemon/services/lda.plist +++ b/daemon/services/lda.plist @@ -17,8 +17,6 @@ {{.Username}} EnvironmentVariables - SHELL - {{.Shell}} HOME {{.Home}} diff --git a/daemon/services/lda.service b/daemon/services/lda.service index cca807f..bde2445 100644 --- a/daemon/services/lda.service +++ b/daemon/services/lda.service @@ -3,7 +3,6 @@ Description=Devzero.io LDA Service After=network.target [Service] -Environment="SHELL={{.Shell}}" Environment="HOME={{.Home}}" User={{.Username}} Group={{.Group}} diff --git a/database/migrations.go b/database/migrations.go index d4e1bfa..b828309 100644 --- a/database/migrations.go +++ b/database/migrations.go @@ -13,6 +13,7 @@ func RunMigrations() { createCommandsTable() createConfigTable() addIndexOnProcesses() + shellTypeToLocation() } func ensureMigrationTableExists() { @@ -110,9 +111,7 @@ func createConfigTable() { home_dir TEXT NOT NULL, lda_dir TEXT NOT NULL, is_root BOOLEAN NOT NULL, - exe_path TEXT NOT NULL, - shell_type INTEGER NOT NULL, - shell_location TEXT NOT NULL + exe_path TEXT NOT NULL );` _, err := DB.Exec(createConfigTableSQL) @@ -124,6 +123,26 @@ func createConfigTable() { } } +func shellTypeToLocation() { + migrationName := "shell_type_to_location" + if !migrationApplied(migrationName) { + shellTypeToLocationSQL := ` + CREATE TABLE IF NOT EXISTS shell_type_to_location ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + shell_type INTEGER NOT NULL, + shell_location TEXT NOT NULL, + config_id INTEGER NOT NULL REFERENCES config(id) + );` + + _, err := DB.Exec(shellTypeToLocationSQL) + if err != nil { + fmt.Fprintf(config.SysConfig.ErrOut, "Failed to add shell_type_to_location column: %s\n", err) + os.Exit(1) + } + recordMigration(migrationName) + } +} + func migrationApplied(migrationName string) bool { var count int err := DB.Get(&count, "SELECT COUNT(*) FROM schema_migrations WHERE migration_name = ?", migrationName) diff --git a/user/user.go b/user/user.go index b9073c7..746079d 100644 --- a/user/user.go +++ b/user/user.go @@ -40,10 +40,8 @@ type Config struct { IsRoot bool `json:"is_root" db:"is_root"` // ExePath is the path to the lda binary ExePath string `json:"exe_path" db:"exe_path"` - // ShellType is the type of the shell - ShellType int64 `json:"shell_type" db:"shell_type"` - // ShellLocation is the location of the shell - ShellLocation string `json:"shell_location" db:"shell_location"` + // ShellTypeToLocation is a map of shell type to location + ShellTypeToLocation map[config.ShellType]string `json:"shell_type_to_location" db:"shell_type_to_location"` // User is the user that executed the command (if sudo) User *user.User `json:"-" db:"-"` } @@ -63,12 +61,38 @@ func GetConfig() (*Config, error) { // InsertConfig inserts Config used to configure the system func InsertConfig(osConfig Config) error { - query := `INSERT INTO config (os, os_name, home_dir, lda_dir, is_root, exe_path, shell_type, shell_location) - VALUES (:os, :os_name, :home_dir, :lda_dir, :is_root, :exe_path, :shell_type, :shell_location)` + query := `INSERT INTO config (os, os_name, home_dir, lda_dir, is_root, exe_path) + VALUES (:os, :os_name, :home_dir, :lda_dir, :is_root, :exe_path)` _, err := database.DB.NamedExec(query, osConfig) + if err != nil { + return err + } - return err + // drop all records in the table + _, err = database.DB.Exec("DELETE FROM shell_type_to_location") + if err != nil { + return err + } + + // get the current config to retrieve the id + currCfg, err := GetConfig() + // should never really happen cuz the config was just inserted + if err != nil { + return err + } + + // osConfig.ShellTypeToLocation is a map of shell type to location + // all the records need to get written to shell_type_to_location table + for shellType, location := range osConfig.ShellTypeToLocation { + // TODO this can be batched + _, err = database.DB.Exec("INSERT INTO shell_type_to_location (shell_type, shell_location, config_id) VALUES (?, ?, ?)", shellType, location, currCfg.Id) + if err != nil { + return err + } + } + + return nil } // UpdateConfig updates an existing Config record in the database @@ -80,18 +104,34 @@ func UpdateConfig(osConfig Config) error { lda_dir = :lda_dir, is_root = :is_root, exe_path = :exe_path, - shell_type = :shell_type, - shell_location = :shell_location WHERE id = :id` _, err := database.DB.NamedExec(query, osConfig) + if err != nil { + return err + } + + // drop all records in the table + _, err = database.DB.Exec("DELETE FROM shell_type_to_location") + if err != nil { + return err + } + + // osConfig.ShellTypeToLocation is a map of shell type to location + // all the records need to get written to shell_type_to_location table + for shellType, location := range osConfig.ShellTypeToLocation { + // TODO this can be batched + _, err = database.DB.Exec("INSERT INTO shell_type_to_location (shell_type, shell_location, config_id) VALUES (?, ?, ?)", shellType, location, osConfig.Id) + if err != nil { + return err + } + } return err } // ConfigureUserSystemInfo configures the user system information and prompts the user to update the configuration if necessary. func ConfigureUserSystemInfo(currentConf *Config) { - // Retrieve the existing configuration from the database. existingConf, err := GetConfig() if err != nil && err != sql.ErrNoRows { @@ -134,22 +174,24 @@ func ConfigureUserSystemInfo(currentConf *Config) { } if result == YesUpdate { - shellType, shellLocation, err := config.GetShell() - if err != nil { - logging.Log.Error().Err(err).Msg("Failed to setup shell") - os.Exit(1) - } - - currentConf.ShellType = int64(shellType) - currentConf.ShellLocation = shellLocation - - currentConf.Id = existingConf.Id - if err := UpdateConfig(*currentConf); err != nil { - logging.Log.Error().Err(err).Msg("Failed to update configuration") - fmt.Fprintf(config.SysConfig.ErrOut, "Failed to update configuration: %s\n", err) - os.Exit(1) + // if shell config is already set by however the binary was invoked, lets respect it + // if not, lets set it up + if len(currentConf.ShellTypeToLocation) == 0 { + shellTypeToLocation, err := config.GetShell() + if err != nil { + logging.Log.Error().Err(err).Msg("Failed to setup shell") + os.Exit(1) + } + currentConf.ShellTypeToLocation = shellTypeToLocation + + currentConf.Id = existingConf.Id + if err := UpdateConfig(*currentConf); err != nil { + logging.Log.Error().Err(err).Msg("Failed to update configuration") + fmt.Fprintf(config.SysConfig.ErrOut, "Failed to update configuration: %s\n", err) + os.Exit(1) + } + logging.Log.Info().Msg("Configuration updated to current settings.") } - logging.Log.Info().Msg("Configuration updated to current settings.") Conf = currentConf } else { existingConf.User = currentConf.User @@ -166,22 +208,20 @@ func ConfigureUserSystemInfo(currentConf *Config) { logging.Log.Debug().Msg("No config found, creating new one") - shellType, shellLocation, err := config.GetShell() - if err != nil { - logging.Log.Error().Err(err).Msg("Failed to setup shell") - os.Exit(1) + if len(currentConf.ShellTypeToLocation) == 0 { + shellTypeToLocation, err := config.GetShell() + if err != nil { + logging.Log.Error().Err(err).Msg("Failed to setup shell") + os.Exit(1) + } + currentConf.ShellTypeToLocation = shellTypeToLocation } - - currentConf.ShellType = int64(shellType) - currentConf.ShellLocation = shellLocation - logging.Log.Debug().Msgf("Shell config: %+v", currentConf) if err := InsertConfig(*currentConf); err != nil { fmt.Fprintf(config.SysConfig.ErrOut, "Failed to insert os config: %s\n", err) os.Exit(1) } - logging.Log.Debug().Msg("Config inserted") Conf = currentConf