Skip to content

Commit ce6f060

Browse files
null93null93
committed
Refactoring to reuse picking logic
1 parent c9cfd9e commit ce6f060

File tree

6 files changed

+121
-196
lines changed

6 files changed

+121
-196
lines changed

‎README.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77

88
## About
99

10-
Knox is a powerful utility designed to simplify and streamline the process of managing AWS credentials. Whether you're frequently switching between different AWS profiles or managing credentials issuance in an SSO environment, Knox provides a straightforward CLI tool to handle these tasks effortlessly. Commands like `knox select`, `knox creds select`, `knox creds last-used`, and `knox clean` make it easy to navigate and manipulate your AWS credential configurations. Its configurable nature, showcased in the `~/.aws/config` setup recommendations, ensures seamless integration into your AWS workflows. Whether you're in development, staging, or production, Knox helps maintain efficient and secure AWS credential management.
10+
Knox is a powerful utility designed to simplify and streamline the process of managing AWS credentials. Whether you're frequently switching between different AWS profiles or managing credentials issuance in an SSO environment, Knox provides a straightforward CLI tool to handle these tasks effortlessly. Commands like `knox select`, `knox last-used`, and `knox clean` make it easy to navigate and manipulate your AWS credential configurations. Its configurable nature, showcased in the `~/.aws/config` setup recommendations, ensures seamless integration into your AWS workflows. Whether you're in development, staging, or production, Knox helps maintain efficient and secure AWS credential management.
11+
12+
Additionally the `knox connect` command provides a simple way to start an SSM session with an EC2 instance. This feature is particularly useful for users who frequently SSH into EC2 instances using SSM Session Manager. Knox allows you to switch between different AWS profiles and start an interactive session with a specific instance using a single command.
1113

1214
## Install
1315

‎internal/connect.go

Lines changed: 14 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -15,74 +15,25 @@ const (
1515
SESSION_MANAGER_PLUGIN_URL = "https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html"
1616
)
1717

18-
var (
19-
connectSessionName string
20-
connectAccountId string
21-
connectRoleName string
22-
connectInstanceId string
23-
connectUid uint32 = 0
24-
)
25-
2618
var connectCmd = &cobra.Command{
2719
Use: "connect",
28-
Short: "Connect to a specific EC2 instance using AWS session-manager-plugin",
20+
Short: "Connect to an EC2 instance using session-manager-plugin",
2921
Args: cobra.ExactArgs(0),
3022
Run: func(cmd *cobra.Command, args []string) {
3123
var err error
32-
var sessions credentials.Sessions
33-
var session *credentials.Session
34-
var roles credentials.Roles
3524
var role *credentials.Role
3625
var binaryPath string
37-
if sessions, err = credentials.GetSessions(); err != nil {
38-
ExitWithError(1, "failed to get configured sessions", err)
39-
}
40-
if connectSessionName == "" {
41-
if connectSessionName, err = tui.SelectSession(sessions); err != nil {
42-
ExitWithError(2, "failed to pick an sso session", err)
43-
}
44-
}
45-
if session = sessions.FindByName(connectSessionName); session == nil {
46-
ExitWithError(3, "session with passed name not found", err)
47-
}
48-
if session.ClientToken == nil || session.ClientToken.IsExpired() {
49-
if err = tui.ClientLogin(session); err != nil {
50-
ExitWithError(4, "failed to authorize device login", err)
51-
}
52-
}
53-
if connectAccountId == "" {
54-
if connectAccountId, err = tui.SelectAccount(session); err != nil {
55-
ExitWithError(5, "failed to pick an account id", err)
56-
}
57-
}
58-
if roles, err = session.GetRoles(connectAccountId); err != nil {
59-
ExitWithError(6, "failed to get roles", err)
60-
}
61-
if connectRoleName == "" {
62-
if connectRoleName, err = tui.SelectRole(roles); err != nil {
63-
ExitWithError(7, "failed to pick a role", err)
64-
}
65-
}
66-
if role = roles.FindByName(connectRoleName); role == nil {
67-
ExitWithError(8, "role with passed name not found", err)
68-
}
69-
if role.Credentials == nil || role.Credentials.IsExpired() {
70-
if err = session.RefreshRoleCredentials(role); err != nil {
71-
ExitWithError(9, "failed to get credentials", err)
72-
}
73-
if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil {
74-
ExitWithError(10, "failed to save credentials", err)
75-
}
76-
}
77-
if err := role.MarkLastUsed(); err != nil {
78-
ExitWithError(11, "failed to mark last used role", err)
26+
if view == "session" {
27+
_, _, _, role = SelectRoleCredentialsStartingFromSession()
28+
} else {
29+
_, _, role = SelectRoleCredentialsStartingFromCache()
7930
}
80-
if connectInstanceId == "" {
81-
if connectInstanceId, err = tui.SelectInstance(role); err != nil {
31+
if instanceId == "" {
32+
if instanceId, err = tui.SelectInstance(role); err != nil {
8233
ExitWithError(12, "failed to pick an instance", err)
8334
}
8435
}
85-
details, err := role.StartSession(connectInstanceId, connectUid)
36+
details, err := role.StartSession(instanceId, connectUid)
8637
if err != nil {
8738
ExitWithError(13, "failed to start ssm session", err)
8839
}
@@ -95,7 +46,7 @@ var connectCmd = &cobra.Command{
9546
role.Region,
9647
"StartSession",
9748
"", // No Profile
98-
fmt.Sprintf(`{"Target": "%s"}`, connectInstanceId),
49+
fmt.Sprintf(`{"Target": "%s"}`, instanceId),
9950
fmt.Sprintf("https://ssm.%s.amazonaws.com", role.Region),
10051
)
10152
command.Stdin = os.Stdin
@@ -111,9 +62,10 @@ var connectCmd = &cobra.Command{
11162
func init() {
11263
RootCmd.AddCommand(connectCmd)
11364
connectCmd.Flags().SortFlags = true
114-
connectCmd.Flags().StringVarP(&connectSessionName, "sso-session", "s", connectSessionName, "SSO session name")
115-
connectCmd.Flags().StringVarP(&connectAccountId, "account-id", "a", connectAccountId, "AWS account ID")
116-
connectCmd.Flags().StringVarP(&connectRoleName, "role-name", "r", connectRoleName, "AWS role name")
117-
connectCmd.Flags().StringVarP(&connectInstanceId, "instance-id", "i", connectInstanceId, "EC2 instance ID")
65+
connectCmd.Flags().StringVarP(&sessionName, "sso-session", "s", sessionName, "SSO session name")
66+
connectCmd.Flags().StringVarP(&accountId, "account-id", "a", accountId, "AWS account ID")
67+
connectCmd.Flags().StringVarP(&roleName, "role-name", "r", roleName, "AWS role name")
68+
connectCmd.Flags().StringVarP(&instanceId, "instance-id", "i", instanceId, "EC2 instance ID")
69+
connectCmd.Flags().StringVarP(&view, "view", "v", view, "session or cached")
11870
connectCmd.Flags().Uint32VarP(&connectUid, "uid", "u", connectUid, "UID on instance to 'su' to")
11971
}

‎internal/creds-select.go

Lines changed: 0 additions & 55 deletions
This file was deleted.

‎internal/creds.go

Lines changed: 0 additions & 15 deletions
This file was deleted.

‎internal/root.go

Lines changed: 96 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,33 +7,124 @@ import (
77

88
"github.com/null93/aws-knox/pkg/ansi"
99
"github.com/null93/aws-knox/pkg/color"
10+
"github.com/null93/aws-knox/sdk/credentials"
11+
"github.com/null93/aws-knox/sdk/tui"
1012
"github.com/spf13/cobra"
1113
)
1214

1315
var (
14-
Version = "0.0.0"
15-
Debug = false
16+
Version string = "0.0.0"
17+
debug bool = false
18+
view string = "session"
19+
connectUid uint32 = 0
20+
sessionName string
21+
accountId string
22+
roleName string
23+
instanceId string
1624
)
1725

1826
var RootCmd = &cobra.Command{
1927
Use: "knox",
2028
Version: Version,
21-
Short: "knox helps you manage AWS credentials that are created via AWS SSO",
29+
Short: "knox helps you manage AWS role credentials and connect to EC2 instances",
2230
}
2331

2432
func ExitWithError(code int, message string, err error) {
2533
fmt.Printf("Error: %s\n", message)
26-
if err != nil && Debug {
34+
if err != nil && debug {
2735
fmt.Printf("Debug: %s\n", err.Error())
2836
}
2937
os.Exit(code)
3038
}
3139

40+
func SelectRoleCredentialsStartingFromSession() (credentials.Sessions, *credentials.Session, credentials.Roles, *credentials.Role) {
41+
var err error
42+
var sessions credentials.Sessions
43+
var session *credentials.Session
44+
var roles credentials.Roles
45+
var role *credentials.Role
46+
if sessions, err = credentials.GetSessions(); err != nil {
47+
ExitWithError(1, "failed to get configured sessions", err)
48+
}
49+
if sessionName == "" {
50+
if sessionName, err = tui.SelectSession(sessions); err != nil {
51+
ExitWithError(2, "failed to pick an sso session", err)
52+
}
53+
}
54+
if session = sessions.FindByName(sessionName); session == nil {
55+
ExitWithError(3, "session with passed name not found", err)
56+
}
57+
if session.ClientToken == nil || session.ClientToken.IsExpired() {
58+
if err = tui.ClientLogin(session); err != nil {
59+
ExitWithError(4, "failed to authorize device login", err)
60+
}
61+
}
62+
if accountId == "" {
63+
if accountId, err = tui.SelectAccount(session); err != nil {
64+
ExitWithError(5, "failed to pick an account id", err)
65+
}
66+
}
67+
if roles, err = session.GetRoles(accountId); err != nil {
68+
ExitWithError(6, "failed to get roles", err)
69+
}
70+
if roleName == "" {
71+
if roleName, err = tui.SelectRole(roles); err != nil {
72+
ExitWithError(7, "failed to pick a role", err)
73+
}
74+
}
75+
if role = roles.FindByName(roleName); role == nil {
76+
ExitWithError(8, "role with passed name not found", err)
77+
}
78+
if role.Credentials == nil || role.Credentials.IsExpired() {
79+
if err = session.RefreshRoleCredentials(role); err != nil {
80+
ExitWithError(9, "failed to get credentials", err)
81+
}
82+
if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil {
83+
ExitWithError(10, "failed to save credentials", err)
84+
}
85+
}
86+
if err := role.MarkLastUsed(); err != nil {
87+
ExitWithError(11, "failed to mark last used role", err)
88+
}
89+
return sessions, session, roles, role
90+
}
91+
92+
func SelectRoleCredentialsStartingFromCache() (credentials.Sessions, *credentials.Session, *credentials.Role) {
93+
var err error
94+
var sessions credentials.Sessions
95+
var session *credentials.Session
96+
var role *credentials.Role
97+
role, err = tui.SelectRolesCredentials()
98+
if role.Credentials == nil || role.Credentials.IsExpired() {
99+
if sessions, err = credentials.GetSessions(); err != nil {
100+
ExitWithError(1, "failed to parse sso sessions", err)
101+
}
102+
if session = sessions.FindByName(role.SessionName); session == nil {
103+
ExitWithError(2, "failed to find sso session "+role.SessionName, err)
104+
}
105+
if session.ClientToken == nil || session.ClientToken.IsExpired() {
106+
if err = tui.ClientLogin(session); err != nil {
107+
ExitWithError(3, "failed to authorize device login", err)
108+
}
109+
}
110+
if err = session.RefreshRoleCredentials(role); err != nil {
111+
ExitWithError(9, "failed to get credentials", err)
112+
}
113+
if err = role.Credentials.Save(session.Name, role.CacheKey()); err != nil {
114+
ExitWithError(10, "failed to save credentials", err)
115+
}
116+
}
117+
if err = role.MarkLastUsed(); err != nil {
118+
ExitWithError(11, "failed to mark last used role", err)
119+
}
120+
return sessions, session, role
121+
}
122+
32123
func init() {
33124
RootCmd.Flags().SortFlags = true
34125
RootCmd.Root().CompletionOptions.DisableDefaultCmd = true
35126
RootCmd.SetHelpCommand(&cobra.Command{Hidden: true})
36-
RootCmd.PersistentFlags().BoolVarP(&Debug, "debug", "d", Debug, "Debug mode")
127+
RootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", debug, "Debug mode")
37128

38129
if tty, err := os.OpenFile("/dev/tty", syscall.O_WRONLY, 0); err == nil {
39130
ansi.Writer = tty

0 commit comments

Comments
 (0)