From 019cf86c5b315208ac3cd44dea86cdc27c3ad1b5 Mon Sep 17 00:00:00 2001 From: Rafael Grigorian Date: Thu, 5 Sep 2024 23:02:20 -0500 Subject: [PATCH] Added --last-used for connect and sync commands, also accepting one arg for instance search term --- internal/connect.go | 74 ++++++++++++++++++++++++++++++++------------ internal/root.go | 1 + internal/sync.go | 74 ++++++++++++++++++++++++++++++++------------ sdk/picker/picker.go | 4 ++- sdk/tui/tui.go | 12 +++---- 5 files changed, 120 insertions(+), 45 deletions(-) diff --git a/internal/connect.go b/internal/connect.go index f879be2..5353a61 100644 --- a/internal/connect.go +++ b/internal/connect.go @@ -12,36 +12,71 @@ import ( ) var connectCmd = &cobra.Command{ - Use: "connect", + Use: "connect ", Short: "Connect to an EC2 instance using session-manager-plugin", - Args: cobra.ExactArgs(0), + Args: cobra.MaximumNArgs(1), Run: func(cmd *cobra.Command, args []string) { + searchTerm := "" + if len(args) > 0 { + searchTerm = args[0] + } var err error var role *credentials.Role var action string var binaryPath string - for { - if !selectCachedFirst { - action, role = SelectRoleCredentialsStartingFromSession() - } else { - action, role = SelectRoleCredentialsStartingFromCache() - } - if action == "toggle-view" { - toggleView() - continue + if lastUsed { + var err error + var sessions credentials.Sessions + var session *credentials.Session + var roleTemp credentials.Role + if roleTemp, err = credentials.GetLastUsedRole(); err != nil { + ExitWithError(1, "failed to get last used role", err) } - if action == "back" { - goBack() - continue + role = &roleTemp + if role.Credentials == nil || role.Credentials.IsExpired() { + if sessions, err = credentials.GetSessions(); err != nil { + ExitWithError(2, "failed to parse sso sessions", err) + } + if session = sessions.FindByName(role.SessionName); session == nil { + ExitWithError(3, "failed to find sso session "+role.SessionName, err) + } + if session.ClientToken == nil || session.ClientToken.IsExpired() { + if err = tui.ClientLogin(session); err != nil { + ExitWithError(4, "failed to authorize device login", err) + } + } + if err = session.RefreshRoleCredentials(role); err != nil { + ExitWithError(5, "failed to get credentials", err) + } + if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil { + ExitWithError(6, "failed to save credentials", err) + } } - if action == "delete" { - if role != nil && role.Credentials != nil { - role.Credentials.DeleteCache(role.SessionName, role.CacheKey()) + } + for { + if role == nil { + if !selectCachedFirst { + action, role = SelectRoleCredentialsStartingFromSession() + } else { + action, role = SelectRoleCredentialsStartingFromCache() + } + if action == "toggle-view" { + toggleView() + continue + } + if action == "back" { + goBack() + continue + } + if action == "delete" { + if role != nil && role.Credentials != nil { + role.Credentials.DeleteCache(role.SessionName, role.CacheKey()) + } + continue } - continue } if instanceId == "" { - if instanceId, action, err = tui.SelectInstance(role); err != nil { + if instanceId, action, err = tui.SelectInstance(role, searchTerm); err != nil { ExitWithError(19, "failed to pick an instance", err) } else if action == "back" { goBack() @@ -84,5 +119,6 @@ func init() { connectCmd.Flags().StringVarP(&roleName, "role-name", "r", roleName, "AWS role name") connectCmd.Flags().StringVarP(&instanceId, "instance-id", "i", instanceId, "EC2 instance ID") connectCmd.Flags().BoolVarP(&selectCachedFirst, "cached", "c", selectCachedFirst, "select from cached credentials") + connectCmd.Flags().BoolVarP(&lastUsed, "last-used", "l", lastUsed, "select last used credentials") connectCmd.Flags().Uint32VarP(&connectUid, "uid", "u", connectUid, "UID on instance to 'su' to") } diff --git a/internal/root.go b/internal/root.go index 669098f..6c54ea0 100644 --- a/internal/root.go +++ b/internal/root.go @@ -18,6 +18,7 @@ var ( debug bool = false selectCachedFirst bool = false connectUid uint32 = 0 + lastUsed bool = false sessionName string accountId string roleName string diff --git a/internal/sync.go b/internal/sync.go index 6655609..4ca2df3 100644 --- a/internal/sync.go +++ b/internal/sync.go @@ -162,35 +162,70 @@ func rsyncPortForward(role *credentials.Role, instanceId string) { } var syncCmd = &cobra.Command{ - Use: "sync", + Use: "sync ", Short: "start rsyncd and port forward to it", - Args: cobra.ExactArgs(0), + Args: cobra.MaximumNArgs(1), Run: func(cmd *cobra.Command, args []string) { + searchTerm := "" + if len(args) > 0 { + searchTerm = args[0] + } var err error var role *credentials.Role var action string - for { - if !selectCachedFirst { - action, role = SelectRoleCredentialsStartingFromSession() - } else { - action, role = SelectRoleCredentialsStartingFromCache() - } - if action == "toggle-view" { - toggleView() - continue + if lastUsed { + var err error + var sessions credentials.Sessions + var session *credentials.Session + var roleTemp credentials.Role + if roleTemp, err = credentials.GetLastUsedRole(); err != nil { + ExitWithError(1, "failed to get last used role", err) } - if action == "back" { - goBack() - continue + role = &roleTemp + if role.Credentials == nil || role.Credentials.IsExpired() { + if sessions, err = credentials.GetSessions(); err != nil { + ExitWithError(2, "failed to parse sso sessions", err) + } + if session = sessions.FindByName(role.SessionName); session == nil { + ExitWithError(3, "failed to find sso session "+role.SessionName, err) + } + if session.ClientToken == nil || session.ClientToken.IsExpired() { + if err = tui.ClientLogin(session); err != nil { + ExitWithError(4, "failed to authorize device login", err) + } + } + if err = session.RefreshRoleCredentials(role); err != nil { + ExitWithError(5, "failed to get credentials", err) + } + if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil { + ExitWithError(6, "failed to save credentials", err) + } } - if action == "delete" { - if role != nil && role.Credentials != nil { - role.Credentials.DeleteCache(role.SessionName, role.CacheKey()) + } + for { + if role == nil { + if !selectCachedFirst { + action, role = SelectRoleCredentialsStartingFromSession() + } else { + action, role = SelectRoleCredentialsStartingFromCache() + } + if action == "toggle-view" { + toggleView() + continue + } + if action == "back" { + goBack() + continue + } + if action == "delete" { + if role != nil && role.Credentials != nil { + role.Credentials.DeleteCache(role.SessionName, role.CacheKey()) + } + continue } - continue } if instanceId == "" { - if instanceId, action, err = tui.SelectInstance(role); err != nil { + if instanceId, action, err = tui.SelectInstance(role, searchTerm); err != nil { ExitWithError(19, "failed to pick an instance", err) } else if action == "back" { goBack() @@ -224,5 +259,6 @@ func init() { syncCmd.Flags().StringVarP(&instanceId, "instance-id", "i", instanceId, "EC2 instance ID") syncCmd.Flags().Uint16VarP(&rsyncPort, "rsync-port", "P", rsyncPort, "rsync port") syncCmd.Flags().Uint16VarP(&localPort, "local-port", "p", localPort, "local port") + syncCmd.Flags().BoolVarP(&lastUsed, "last-used", "l", lastUsed, "select last used credentials") syncCmd.Flags().BoolVarP(&selectCachedFirst, "cached", "c", selectCachedFirst, "select from cached credentials") } diff --git a/sdk/picker/picker.go b/sdk/picker/picker.go index 9f48935..c0e7e5b 100644 --- a/sdk/picker/picker.go +++ b/sdk/picker/picker.go @@ -186,7 +186,9 @@ func (p *picker) render() { ansi.MoveCursorUp(6 + lines) } -func (p *picker) Pick() (*option, *keys.KeyCode) { +func (p *picker) Pick(initialFilter string) (*option, *keys.KeyCode) { + p.term = initialFilter + p.filter() ansi.HideCursor() defer ansi.ClearDown() defer ansi.ShowCursor() diff --git a/sdk/tui/tui.go b/sdk/tui/tui.go index 5c2485f..4678053 100644 --- a/sdk/tui/tui.go +++ b/sdk/tui/tui.go @@ -84,7 +84,7 @@ func SelectSession(sessions credentials.Sessions) (string, string, error) { } p.AddOption(session.Name, session.Name, session.Region, session.StartUrl, expires) } - selection, firedKeyCode := p.Pick() + selection, firedKeyCode := p.Pick("") if firedKeyCode != nil && *firedKeyCode == keys.Tab { return "", "toggle-view", nil } @@ -114,7 +114,7 @@ func SelectAccount(session *credentials.Session, accountAliases map[string]strin } p.AddOption(account.Id, account.Id, name, account.Email) } - selection, firedKeyCode := p.Pick() + selection, firedKeyCode := p.Pick("") if firedKeyCode != nil && *firedKeyCode == keys.Esc { return "", "back", nil } @@ -139,7 +139,7 @@ func SelectRole(roles credentials.Roles) (string, string, error) { } p.AddOption(role.Name, role.Name, expires) } - selection, firedKeyCode := p.Pick() + selection, firedKeyCode := p.Pick("") if firedKeyCode != nil && *firedKeyCode == keys.Esc { return "", "back", nil } @@ -149,7 +149,7 @@ func SelectRole(roles credentials.Roles) (string, string, error) { return selection.Value.(string), "", nil } -func SelectInstance(role *credentials.Role) (string, string, error) { +func SelectInstance(role *credentials.Role, initialFilter string) (string, string, error) { instances, err := role.GetManagedInstances() if err != nil { return "", "", err @@ -163,7 +163,7 @@ func SelectInstance(role *credentials.Role) (string, string, error) { for _, instance := range instances { p.AddOption(instance.Id, instance.Id, instance.InstanceType, instance.PrivateIpAddress, instance.PublicIpAddress, instance.Name) } - selection, firedKeyCode := p.Pick() + selection, firedKeyCode := p.Pick(initialFilter) if firedKeyCode != nil && *firedKeyCode == keys.Esc { return "", "back", nil } @@ -199,7 +199,7 @@ func SelectRolesCredentials(accountAliases map[string]string) (*credentials.Role } p.AddOption(role, role.SessionName, role.Region, role.AccountId, alias, role.Name, expires) } - selection, firedKeyCode := p.Pick() + selection, firedKeyCode := p.Pick("") if firedKeyCode != nil && *firedKeyCode == keys.Tab { return nil, "toggle-view", nil }