Skip to content

Commit 32acace

Browse files
null93null93
committed
Added go back action and toggle view for connect and select
1 parent ce6f060 commit 32acace

File tree

5 files changed

+195
-84
lines changed

5 files changed

+195
-84
lines changed

‎internal/connect.go

Lines changed: 52 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -19,42 +19,63 @@ var connectCmd = &cobra.Command{
1919
Use: "connect",
2020
Short: "Connect to an EC2 instance using session-manager-plugin",
2121
Args: cobra.ExactArgs(0),
22+
PreRunE: func(cmd *cobra.Command, args []string) error {
23+
if view != "session" && view != "cached" {
24+
return fmt.Errorf("view must be either 'session' or 'cached'")
25+
}
26+
return nil
27+
},
2228
Run: func(cmd *cobra.Command, args []string) {
2329
var err error
2430
var role *credentials.Role
31+
var action string
2532
var binaryPath string
26-
if view == "session" {
27-
_, _, _, role = SelectRoleCredentialsStartingFromSession()
28-
} else {
29-
_, _, role = SelectRoleCredentialsStartingFromCache()
30-
}
31-
if instanceId == "" {
32-
if instanceId, err = tui.SelectInstance(role); err != nil {
33-
ExitWithError(12, "failed to pick an instance", err)
33+
for {
34+
if view == "session" {
35+
action, _, _, _, role = SelectRoleCredentialsStartingFromSession()
36+
} else {
37+
action, _, _, role = SelectRoleCredentialsStartingFromCache()
3438
}
35-
}
36-
details, err := role.StartSession(instanceId, connectUid)
37-
if err != nil {
38-
ExitWithError(13, "failed to start ssm session", err)
39-
}
40-
if binaryPath, err = exec.LookPath("session-manager-plugin"); err != nil {
41-
ExitWithError(14, "failed to find session-manager-plugin, see "+SESSION_MANAGER_PLUGIN_URL, err)
42-
}
43-
command := exec.Command(
44-
binaryPath,
45-
fmt.Sprintf(`{"SessionId": "%s", "TokenValue": "%s", "StreamUrl": "%s"}`, *details.SessionId, *details.TokenValue, *details.StreamUrl),
46-
role.Region,
47-
"StartSession",
48-
"", // No Profile
49-
fmt.Sprintf(`{"Target": "%s"}`, instanceId),
50-
fmt.Sprintf("https://ssm.%s.amazonaws.com", role.Region),
51-
)
52-
command.Stdin = os.Stdin
53-
command.Stdout = os.Stdout
54-
command.Stderr = os.Stderr
55-
command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Foreground: true}
56-
if err = command.Run(); err != nil {
57-
ExitWithError(15, "failed to run session-manager-plugin", err)
39+
if action == "toggle-view" {
40+
toggleView()
41+
continue
42+
}
43+
if action == "back" {
44+
goBack()
45+
continue
46+
}
47+
if instanceId == "" {
48+
if instanceId, action, err = tui.SelectInstance(role); err != nil {
49+
ExitWithError(12, "failed to pick an instance", err)
50+
} else if action == "back" {
51+
goBack()
52+
continue
53+
}
54+
}
55+
details, err := role.StartSession(instanceId, connectUid)
56+
if err != nil {
57+
ExitWithError(13, "failed to start ssm session", err)
58+
}
59+
if binaryPath, err = exec.LookPath("session-manager-plugin"); err != nil {
60+
ExitWithError(14, "failed to find session-manager-plugin, see "+SESSION_MANAGER_PLUGIN_URL, err)
61+
}
62+
command := exec.Command(
63+
binaryPath,
64+
fmt.Sprintf(`{"SessionId": "%s", "TokenValue": "%s", "StreamUrl": "%s"}`, *details.SessionId, *details.TokenValue, *details.StreamUrl),
65+
role.Region,
66+
"StartSession",
67+
"", // No Profile
68+
fmt.Sprintf(`{"Target": "%s"}`, instanceId),
69+
fmt.Sprintf("https://ssm.%s.amazonaws.com", role.Region),
70+
)
71+
command.Stdin = os.Stdin
72+
command.Stdout = os.Stdout
73+
command.Stderr = os.Stderr
74+
command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Foreground: true}
75+
if err = command.Run(); err != nil {
76+
ExitWithError(15, "failed to run session-manager-plugin", err)
77+
}
78+
break
5879
}
5980
},
6081
}

‎internal/root.go

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,32 @@ var RootCmd = &cobra.Command{
2929
Short: "knox helps you manage AWS role credentials and connect to EC2 instances",
3030
}
3131

32+
func toggleView() {
33+
if view == "session" {
34+
view = "cached"
35+
} else {
36+
view = "session"
37+
}
38+
}
39+
40+
func goBack() {
41+
if instanceId != "" {
42+
instanceId = ""
43+
return
44+
}
45+
if roleName != "" {
46+
roleName = ""
47+
return
48+
}
49+
if accountId != "" {
50+
accountId = ""
51+
return
52+
}
53+
if sessionName != "" {
54+
sessionName = ""
55+
}
56+
}
57+
3258
func ExitWithError(code int, message string, err error) {
3359
fmt.Printf("Error: %s\n", message)
3460
if err != nil && debug {
@@ -37,8 +63,9 @@ func ExitWithError(code int, message string, err error) {
3763
os.Exit(code)
3864
}
3965

40-
func SelectRoleCredentialsStartingFromSession() (credentials.Sessions, *credentials.Session, credentials.Roles, *credentials.Role) {
66+
func SelectRoleCredentialsStartingFromSession() (string, credentials.Sessions, *credentials.Session, credentials.Roles, *credentials.Role) {
4167
var err error
68+
var action string
4269
var sessions credentials.Sessions
4370
var session *credentials.Session
4471
var roles credentials.Roles
@@ -47,8 +74,10 @@ func SelectRoleCredentialsStartingFromSession() (credentials.Sessions, *credenti
4774
ExitWithError(1, "failed to get configured sessions", err)
4875
}
4976
if sessionName == "" {
50-
if sessionName, err = tui.SelectSession(sessions); err != nil {
77+
if sessionName, action, err = tui.SelectSession(sessions); err != nil {
5178
ExitWithError(2, "failed to pick an sso session", err)
79+
} else if action != "" {
80+
return action, nil, nil, nil, nil
5281
}
5382
}
5483
if session = sessions.FindByName(sessionName); session == nil {
@@ -60,16 +89,20 @@ func SelectRoleCredentialsStartingFromSession() (credentials.Sessions, *credenti
6089
}
6190
}
6291
if accountId == "" {
63-
if accountId, err = tui.SelectAccount(session); err != nil {
92+
if accountId, action, err = tui.SelectAccount(session); err != nil {
6493
ExitWithError(5, "failed to pick an account id", err)
94+
} else if action != "" {
95+
return action, nil, nil, nil, nil
6596
}
6697
}
6798
if roles, err = session.GetRoles(accountId); err != nil {
6899
ExitWithError(6, "failed to get roles", err)
69100
}
70101
if roleName == "" {
71-
if roleName, err = tui.SelectRole(roles); err != nil {
102+
if roleName, action, err = tui.SelectRole(roles); err != nil {
72103
ExitWithError(7, "failed to pick a role", err)
104+
} else if action != "" {
105+
return action, nil, nil, nil, nil
73106
}
74107
}
75108
if role = roles.FindByName(roleName); role == nil {
@@ -86,15 +119,19 @@ func SelectRoleCredentialsStartingFromSession() (credentials.Sessions, *credenti
86119
if err := role.MarkLastUsed(); err != nil {
87120
ExitWithError(11, "failed to mark last used role", err)
88121
}
89-
return sessions, session, roles, role
122+
return "", sessions, session, roles, role
90123
}
91124

92-
func SelectRoleCredentialsStartingFromCache() (credentials.Sessions, *credentials.Session, *credentials.Role) {
125+
func SelectRoleCredentialsStartingFromCache() (string, credentials.Sessions, *credentials.Session, *credentials.Role) {
93126
var err error
127+
var action string
94128
var sessions credentials.Sessions
95129
var session *credentials.Session
96130
var role *credentials.Role
97-
role, err = tui.SelectRolesCredentials()
131+
role, action, err = tui.SelectRolesCredentials()
132+
if action != "" {
133+
return action, nil, nil, nil
134+
}
98135
if role.Credentials == nil || role.Credentials.IsExpired() {
99136
if sessions, err = credentials.GetSessions(); err != nil {
100137
ExitWithError(1, "failed to parse sso sessions", err)
@@ -117,7 +154,7 @@ func SelectRoleCredentialsStartingFromCache() (credentials.Sessions, *credential
117154
if err = role.MarkLastUsed(); err != nil {
118155
ExitWithError(11, "failed to mark last used role", err)
119156
}
120-
return sessions, session, role
157+
return "", sessions, session, role
121158
}
122159

123160
func init() {

‎internal/select.go

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,27 @@ var selectCmd = &cobra.Command{
1919
},
2020
Run: func(cmd *cobra.Command, args []string) {
2121
var role *credentials.Role
22-
if view == "session" {
23-
_, _, _, role = SelectRoleCredentialsStartingFromSession()
24-
} else {
25-
_, _, role = SelectRoleCredentialsStartingFromCache()
26-
}
27-
if json, err := role.Credentials.ToJSON(); err != nil {
28-
ExitWithError(12, "failed to convert credentials to json", err)
29-
} else {
30-
fmt.Println(json)
22+
var action string
23+
for {
24+
if view == "session" {
25+
action, _, _, _, role = SelectRoleCredentialsStartingFromSession()
26+
} else {
27+
action, _, _, role = SelectRoleCredentialsStartingFromCache()
28+
}
29+
if action == "toggle-view" {
30+
toggleView()
31+
continue
32+
}
33+
if action == "back" {
34+
goBack()
35+
continue
36+
}
37+
if json, err := role.Credentials.ToJSON(); err != nil {
38+
ExitWithError(12, "failed to convert credentials to json", err)
39+
} else {
40+
fmt.Println(json)
41+
break
42+
}
3143
}
3244
},
3345
}

‎sdk/picker/picker.go

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
)
1212

1313
type picker struct {
14+
actions []action
1415
options []option
1516
filtered []*option
1617
selectedIndex int
@@ -29,8 +30,15 @@ type option struct {
2930
Value interface{}
3031
}
3132

33+
type action struct {
34+
key keys.KeyCode
35+
name string
36+
description string
37+
}
38+
3239
func NewPicker() *picker {
3340
p := picker{
41+
actions: []action{},
3442
options: []option{},
3543
filtered: []*option{},
3644
selectedIndex: 0,
@@ -91,6 +99,10 @@ func (p *picker) AddOption(value interface{}, cols ...string) {
9199
p.filtered = append(p.filtered, &o)
92100
}
93101

102+
func (p *picker) AddAction(key keys.KeyCode, name string, description string) {
103+
p.actions = append(p.actions, action{key, name, description})
104+
}
105+
94106
func (p *picker) filter() {
95107
p.filtered = []*option{}
96108
p.selectedIndex = 0
@@ -154,14 +166,15 @@ func (p *picker) render() {
154166
} else {
155167
DefaultStyle.Printfln("")
156168
}
157-
DefaultStyle.Printf(
158-
darkGray(" %d/%d items •", len(p.filtered), len(p.options)) +
159-
lightGray(" ↑ ") + darkGray("up •") +
160-
lightGray(" ↓ ") + darkGray("down •") +
161-
lightGray(" enter ") + darkGray("choose •") +
162-
lightGray(" esc ") + darkGray("quit") +
163-
color.ResetStyle,
164-
)
169+
helpMenu := darkGray(" %d/%d items •", len(p.filtered), len(p.options))
170+
helpMenu += lightGray(" ↑ ") + darkGray("up •")
171+
helpMenu += lightGray(" ↓ ") + darkGray("down •")
172+
helpMenu += lightGray(" enter ") + darkGray("choose •")
173+
for _, action := range p.actions {
174+
helpMenu += lightGray(" %s ", action.name) + darkGray("%s •", action.description)
175+
}
176+
helpMenu += lightGray(" ctl+c ") + darkGray("quit") + color.ResetStyle
177+
DefaultStyle.Printf(helpMenu)
165178
DefaultStyle.Printfln("")
166179
lines := len(p.filtered)
167180
if p.maxHeight < lines {
@@ -170,13 +183,20 @@ func (p *picker) render() {
170183
ansi.MoveCursorUp(7 + lines)
171184
}
172185

173-
func (p *picker) Pick() *option {
186+
func (p *picker) Pick() (*option, *keys.KeyCode) {
174187
ansi.HideCursor()
175188
defer ansi.ClearDown()
176189
defer ansi.ShowCursor()
177190
p.render()
191+
var firedKeyCode keys.KeyCode
178192
keyboard.Listen(func(key keys.Key) (stop bool, err error) {
179-
if key.Code == keys.CtrlC || key.Code == keys.Escape {
193+
for _, action := range p.actions {
194+
if key.Code == action.key {
195+
firedKeyCode = action.key
196+
return true, nil
197+
}
198+
}
199+
if key.Code == keys.CtrlC {
180200
p.selectedIndex = -1
181201
return true, nil
182202
}
@@ -220,10 +240,10 @@ func (p *picker) Pick() *option {
220240
return false, nil
221241
})
222242
if p.selectedIndex < 0 {
223-
return nil
243+
return nil, nil
224244
}
225245
if p.selectedIndex >= len(p.filtered) {
226-
return nil
246+
return nil, nil
227247
}
228-
return p.filtered[p.selectedIndex]
248+
return p.filtered[p.selectedIndex], &firedKeyCode
229249
}

0 commit comments

Comments
 (0)