Skip to content

Commit 019cf86

Browse files
null93null93
committed
Added --last-used for connect and sync commands, also accepting one arg for instance search term
1 parent 4a6b0b2 commit 019cf86

File tree

5 files changed

+120
-45
lines changed

5 files changed

+120
-45
lines changed

‎internal/connect.go

Lines changed: 55 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -12,36 +12,71 @@ import (
1212
)
1313

1414
var connectCmd = &cobra.Command{
15-
Use: "connect",
15+
Use: "connect <instance-search-term>",
1616
Short: "Connect to an EC2 instance using session-manager-plugin",
17-
Args: cobra.ExactArgs(0),
17+
Args: cobra.MaximumNArgs(1),
1818
Run: func(cmd *cobra.Command, args []string) {
19+
searchTerm := ""
20+
if len(args) > 0 {
21+
searchTerm = args[0]
22+
}
1923
var err error
2024
var role *credentials.Role
2125
var action string
2226
var binaryPath string
23-
for {
24-
if !selectCachedFirst {
25-
action, role = SelectRoleCredentialsStartingFromSession()
26-
} else {
27-
action, role = SelectRoleCredentialsStartingFromCache()
28-
}
29-
if action == "toggle-view" {
30-
toggleView()
31-
continue
27+
if lastUsed {
28+
var err error
29+
var sessions credentials.Sessions
30+
var session *credentials.Session
31+
var roleTemp credentials.Role
32+
if roleTemp, err = credentials.GetLastUsedRole(); err != nil {
33+
ExitWithError(1, "failed to get last used role", err)
3234
}
33-
if action == "back" {
34-
goBack()
35-
continue
35+
role = &roleTemp
36+
if role.Credentials == nil || role.Credentials.IsExpired() {
37+
if sessions, err = credentials.GetSessions(); err != nil {
38+
ExitWithError(2, "failed to parse sso sessions", err)
39+
}
40+
if session = sessions.FindByName(role.SessionName); session == nil {
41+
ExitWithError(3, "failed to find sso session "+role.SessionName, err)
42+
}
43+
if session.ClientToken == nil || session.ClientToken.IsExpired() {
44+
if err = tui.ClientLogin(session); err != nil {
45+
ExitWithError(4, "failed to authorize device login", err)
46+
}
47+
}
48+
if err = session.RefreshRoleCredentials(role); err != nil {
49+
ExitWithError(5, "failed to get credentials", err)
50+
}
51+
if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil {
52+
ExitWithError(6, "failed to save credentials", err)
53+
}
3654
}
37-
if action == "delete" {
38-
if role != nil && role.Credentials != nil {
39-
role.Credentials.DeleteCache(role.SessionName, role.CacheKey())
55+
}
56+
for {
57+
if role == nil {
58+
if !selectCachedFirst {
59+
action, role = SelectRoleCredentialsStartingFromSession()
60+
} else {
61+
action, role = SelectRoleCredentialsStartingFromCache()
62+
}
63+
if action == "toggle-view" {
64+
toggleView()
65+
continue
66+
}
67+
if action == "back" {
68+
goBack()
69+
continue
70+
}
71+
if action == "delete" {
72+
if role != nil && role.Credentials != nil {
73+
role.Credentials.DeleteCache(role.SessionName, role.CacheKey())
74+
}
75+
continue
4076
}
41-
continue
4277
}
4378
if instanceId == "" {
44-
if instanceId, action, err = tui.SelectInstance(role); err != nil {
79+
if instanceId, action, err = tui.SelectInstance(role, searchTerm); err != nil {
4580
ExitWithError(19, "failed to pick an instance", err)
4681
} else if action == "back" {
4782
goBack()
@@ -84,5 +119,6 @@ func init() {
84119
connectCmd.Flags().StringVarP(&roleName, "role-name", "r", roleName, "AWS role name")
85120
connectCmd.Flags().StringVarP(&instanceId, "instance-id", "i", instanceId, "EC2 instance ID")
86121
connectCmd.Flags().BoolVarP(&selectCachedFirst, "cached", "c", selectCachedFirst, "select from cached credentials")
122+
connectCmd.Flags().BoolVarP(&lastUsed, "last-used", "l", lastUsed, "select last used credentials")
87123
connectCmd.Flags().Uint32VarP(&connectUid, "uid", "u", connectUid, "UID on instance to 'su' to")
88124
}

‎internal/root.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ var (
1818
debug bool = false
1919
selectCachedFirst bool = false
2020
connectUid uint32 = 0
21+
lastUsed bool = false
2122
sessionName string
2223
accountId string
2324
roleName string

‎internal/sync.go

Lines changed: 55 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -162,35 +162,70 @@ func rsyncPortForward(role *credentials.Role, instanceId string) {
162162
}
163163

164164
var syncCmd = &cobra.Command{
165-
Use: "sync",
165+
Use: "sync <instance-search-term>",
166166
Short: "start rsyncd and port forward to it",
167-
Args: cobra.ExactArgs(0),
167+
Args: cobra.MaximumNArgs(1),
168168
Run: func(cmd *cobra.Command, args []string) {
169+
searchTerm := ""
170+
if len(args) > 0 {
171+
searchTerm = args[0]
172+
}
169173
var err error
170174
var role *credentials.Role
171175
var action string
172-
for {
173-
if !selectCachedFirst {
174-
action, role = SelectRoleCredentialsStartingFromSession()
175-
} else {
176-
action, role = SelectRoleCredentialsStartingFromCache()
177-
}
178-
if action == "toggle-view" {
179-
toggleView()
180-
continue
176+
if lastUsed {
177+
var err error
178+
var sessions credentials.Sessions
179+
var session *credentials.Session
180+
var roleTemp credentials.Role
181+
if roleTemp, err = credentials.GetLastUsedRole(); err != nil {
182+
ExitWithError(1, "failed to get last used role", err)
181183
}
182-
if action == "back" {
183-
goBack()
184-
continue
184+
role = &roleTemp
185+
if role.Credentials == nil || role.Credentials.IsExpired() {
186+
if sessions, err = credentials.GetSessions(); err != nil {
187+
ExitWithError(2, "failed to parse sso sessions", err)
188+
}
189+
if session = sessions.FindByName(role.SessionName); session == nil {
190+
ExitWithError(3, "failed to find sso session "+role.SessionName, err)
191+
}
192+
if session.ClientToken == nil || session.ClientToken.IsExpired() {
193+
if err = tui.ClientLogin(session); err != nil {
194+
ExitWithError(4, "failed to authorize device login", err)
195+
}
196+
}
197+
if err = session.RefreshRoleCredentials(role); err != nil {
198+
ExitWithError(5, "failed to get credentials", err)
199+
}
200+
if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil {
201+
ExitWithError(6, "failed to save credentials", err)
202+
}
185203
}
186-
if action == "delete" {
187-
if role != nil && role.Credentials != nil {
188-
role.Credentials.DeleteCache(role.SessionName, role.CacheKey())
204+
}
205+
for {
206+
if role == nil {
207+
if !selectCachedFirst {
208+
action, role = SelectRoleCredentialsStartingFromSession()
209+
} else {
210+
action, role = SelectRoleCredentialsStartingFromCache()
211+
}
212+
if action == "toggle-view" {
213+
toggleView()
214+
continue
215+
}
216+
if action == "back" {
217+
goBack()
218+
continue
219+
}
220+
if action == "delete" {
221+
if role != nil && role.Credentials != nil {
222+
role.Credentials.DeleteCache(role.SessionName, role.CacheKey())
223+
}
224+
continue
189225
}
190-
continue
191226
}
192227
if instanceId == "" {
193-
if instanceId, action, err = tui.SelectInstance(role); err != nil {
228+
if instanceId, action, err = tui.SelectInstance(role, searchTerm); err != nil {
194229
ExitWithError(19, "failed to pick an instance", err)
195230
} else if action == "back" {
196231
goBack()
@@ -224,5 +259,6 @@ func init() {
224259
syncCmd.Flags().StringVarP(&instanceId, "instance-id", "i", instanceId, "EC2 instance ID")
225260
syncCmd.Flags().Uint16VarP(&rsyncPort, "rsync-port", "P", rsyncPort, "rsync port")
226261
syncCmd.Flags().Uint16VarP(&localPort, "local-port", "p", localPort, "local port")
262+
syncCmd.Flags().BoolVarP(&lastUsed, "last-used", "l", lastUsed, "select last used credentials")
227263
syncCmd.Flags().BoolVarP(&selectCachedFirst, "cached", "c", selectCachedFirst, "select from cached credentials")
228264
}

‎sdk/picker/picker.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,9 @@ func (p *picker) render() {
186186
ansi.MoveCursorUp(6 + lines)
187187
}
188188

189-
func (p *picker) Pick() (*option, *keys.KeyCode) {
189+
func (p *picker) Pick(initialFilter string) (*option, *keys.KeyCode) {
190+
p.term = initialFilter
191+
p.filter()
190192
ansi.HideCursor()
191193
defer ansi.ClearDown()
192194
defer ansi.ShowCursor()

‎sdk/tui/tui.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func SelectSession(sessions credentials.Sessions) (string, string, error) {
8484
}
8585
p.AddOption(session.Name, session.Name, session.Region, session.StartUrl, expires)
8686
}
87-
selection, firedKeyCode := p.Pick()
87+
selection, firedKeyCode := p.Pick("")
8888
if firedKeyCode != nil && *firedKeyCode == keys.Tab {
8989
return "", "toggle-view", nil
9090
}
@@ -114,7 +114,7 @@ func SelectAccount(session *credentials.Session, accountAliases map[string]strin
114114
}
115115
p.AddOption(account.Id, account.Id, name, account.Email)
116116
}
117-
selection, firedKeyCode := p.Pick()
117+
selection, firedKeyCode := p.Pick("")
118118
if firedKeyCode != nil && *firedKeyCode == keys.Esc {
119119
return "", "back", nil
120120
}
@@ -139,7 +139,7 @@ func SelectRole(roles credentials.Roles) (string, string, error) {
139139
}
140140
p.AddOption(role.Name, role.Name, expires)
141141
}
142-
selection, firedKeyCode := p.Pick()
142+
selection, firedKeyCode := p.Pick("")
143143
if firedKeyCode != nil && *firedKeyCode == keys.Esc {
144144
return "", "back", nil
145145
}
@@ -149,7 +149,7 @@ func SelectRole(roles credentials.Roles) (string, string, error) {
149149
return selection.Value.(string), "", nil
150150
}
151151

152-
func SelectInstance(role *credentials.Role) (string, string, error) {
152+
func SelectInstance(role *credentials.Role, initialFilter string) (string, string, error) {
153153
instances, err := role.GetManagedInstances()
154154
if err != nil {
155155
return "", "", err
@@ -163,7 +163,7 @@ func SelectInstance(role *credentials.Role) (string, string, error) {
163163
for _, instance := range instances {
164164
p.AddOption(instance.Id, instance.Id, instance.InstanceType, instance.PrivateIpAddress, instance.PublicIpAddress, instance.Name)
165165
}
166-
selection, firedKeyCode := p.Pick()
166+
selection, firedKeyCode := p.Pick(initialFilter)
167167
if firedKeyCode != nil && *firedKeyCode == keys.Esc {
168168
return "", "back", nil
169169
}
@@ -199,7 +199,7 @@ func SelectRolesCredentials(accountAliases map[string]string) (*credentials.Role
199199
}
200200
p.AddOption(role, role.SessionName, role.Region, role.AccountId, alias, role.Name, expires)
201201
}
202-
selection, firedKeyCode := p.Pick()
202+
selection, firedKeyCode := p.Pick("")
203203
if firedKeyCode != nil && *firedKeyCode == keys.Tab {
204204
return nil, "toggle-view", nil
205205
}

0 commit comments

Comments
 (0)