6 files changed, 508 insertions, 0 deletions
diff --git a/internal/command/commandargs/authorized_keys.go b/internal/command/commandargs/authorized_keys.go
new file mode 100644
index 0000000..2733954
--- /dev/null
+++ b/internal/command/commandargs/authorized_keys.go
@@ -0,0 +1,51 @@
+package commandargs
+
+import (
+	"errors"
+	"fmt"
+)
+
+type AuthorizedKeys struct {
+	Arguments    []string
+	ExpectedUser string
+	ActualUser   string
+	Key          string
+}
+
+func (ak *AuthorizedKeys) Parse() error {
+	if err := ak.validate(); err != nil {
+		return err
+	}
+
+	ak.ExpectedUser = ak.Arguments[0]
+	ak.ActualUser = ak.Arguments[1]
+	ak.Key = ak.Arguments[2]
+
+	return nil
+}
+
+func (ak *AuthorizedKeys) GetArguments() []string {
+	return ak.Arguments
+}
+
+func (ak *AuthorizedKeys) validate() error {
+	argsSize := len(ak.Arguments)
+
+	if argsSize != 3 {
+		return errors.New(fmt.Sprintf("# Insufficient arguments. %d. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>", argsSize))
+	}
+
+	expectedUsername := ak.Arguments[0]
+	actualUsername := ak.Arguments[1]
+	key := ak.Arguments[2]
+
+	if expectedUsername == "" || actualUsername == "" {
+		return errors.New("# No username provided")
+	}
+
+	if key == "" {
+		return errors.New("# No key provided")
+	}
+
+	return nil
+}
diff --git a/internal/command/commandargs/authorized_principals.go b/internal/command/commandargs/authorized_principals.go
new file mode 100644
index 0000000..746ae3f
--- /dev/null
+++ b/internal/command/commandargs/authorized_principals.go
@@ -0,0 +1,50 @@
+package commandargs
+
+import (
+	"errors"
+	"fmt"
+)
+
+type AuthorizedPrincipals struct {
+	Arguments  []string
+	KeyId      string
+	Principals []string
+}
+
+func (ap *AuthorizedPrincipals) Parse() error {
+	if err := ap.validate(); err != nil {
+		return err
+	}
+
+	ap.KeyId = ap.Arguments[0]
+	ap.Principals = ap.Arguments[1:]
+
+	return nil
+}
+
+func (ap *AuthorizedPrincipals) GetArguments() []string {
+	return ap.Arguments
+}
+
+func (ap *AuthorizedPrincipals) validate() error {
+	argsSize := len(ap.Arguments)
+
+	if argsSize < 2 {
+		return errors.New(fmt.Sprintf("# Insufficient arguments. %d. Usage\n#\tgitlab-shell-authorized-principals-check <key-id> <principal1> [<principal2>...]", argsSize))
+	}
+
+	keyId := ap.Arguments[0]
+	principals := ap.Arguments[1:]
+
+	if keyId == "" {
+		return errors.New("# No key_id provided")
+	}
+
+	for _, principal := range principals {
+		if principal == "" {
+			return errors.New("# An invalid principal was provided")
+		}
+	}
+
+	return nil
+}
diff --git a/internal/command/commandargs/command_args.go b/internal/command/commandargs/command_args.go
new file mode 100644
index 0000000..4831134
--- /dev/null
+++ b/internal/command/commandargs/command_args.go
@@ -0,0 +1,31 @@
+package commandargs
+
+import (
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
+)
+
+type CommandType string
+
+type CommandArgs interface {
+	Parse() error
+	GetArguments() []string
+}
+
+func Parse(e *executable.Executable, arguments []string) (CommandArgs, error) {
+	var args CommandArgs = &GenericArgs{Arguments: arguments}
+
+	switch e.Name {
+	case executable.GitlabShell:
+		args = &Shell{Arguments: arguments}
+	case executable.AuthorizedKeysCheck:
+		args = &AuthorizedKeys{Arguments: arguments}
+	case executable.AuthorizedPrincipalsCheck:
+		args = &AuthorizedPrincipals{Arguments: arguments}
+	}
+
+	if err := args.Parse(); err != nil {
+		return nil, err
+	}
+
+	return args, nil
+}
diff --git a/internal/command/commandargs/command_args_test.go b/internal/command/commandargs/command_args_test.go
new file mode 100644
index 0000000..9f1575d
--- /dev/null
+++ b/internal/command/commandargs/command_args_test.go
@@ -0,0 +1,231 @@
+package commandargs
+
+import (
+	"testing"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseSuccess(t *testing.T) {
+	testCases := []struct {
+		desc         string
+		executable   *executable.Executable
+		environment  map[string]string
+		arguments    []string
+		expectedArgs CommandArgs
+	}{
+		// Setting the used env variables for every case to ensure we're
+		// not using anything set in the original env.
+		{
+			desc:       "It sets discover as the command when the command string was empty",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{}, CommandType: Discover},
+		},
+		{
+			desc:       "It finds the key id in any passed arguments",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "",
+			},
+			arguments:    []string{"hello", "key-123"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "key-123"}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+		}, {
+			desc:       "It finds the username in any passed arguments",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "",
+			},
+			arguments:    []string{"hello", "username-jane-doe"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "username-jane-doe"}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+		}, {
+			desc:       "It parses 2fa_recovery_codes command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
+		}, {
+			desc:       "It parses git-receive-pack command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git-receive-pack group/repo",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc:       "It parses git-receive-pack command and a project with single quotes",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git receive-pack 'group/repo'",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc:       `It parses "git receive-pack" command`,
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git receive-pack "group/repo"`,
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc:       `It parses a command followed by control characters`,
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git-receive-pack group/repo; any command`,
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc:       "It parses git-upload-pack command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git upload-pack "group/repo"`,
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
+		}, {
+			desc:       "It parses git-upload-archive command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git-upload-archive 'group/repo'",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
+		}, {
+			desc:       "It parses git-lfs-authenticate command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate 'group/repo' download",
+			},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+		}, {
+			desc:         "It parses authorized-keys command",
+			executable:   &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:    []string{"git", "git", "key"},
+			expectedArgs: &AuthorizedKeys{Arguments: []string{"git", "git", "key"}, ExpectedUser: "git", ActualUser: "git", Key: "key"},
+		}, {
+			desc:         "It parses authorized-principals command",
+			executable:   &executable.Executable{Name: executable.AuthorizedPrincipalsCheck},
+			arguments:    []string{"key", "principal-1", "principal-2"},
+			expectedArgs: &AuthorizedPrincipals{Arguments: []string{"key", "principal-1", "principal-2"}, KeyId: "key", Principals: []string{"principal-1", "principal-2"}},
+		}, {
+			desc:         "Unknown executable",
+			executable:   &executable.Executable{Name: "unknown"},
+			arguments:    []string{},
+			expectedArgs: &GenericArgs{Arguments: []string{}},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
+
+			result, err := Parse(tc.executable, tc.arguments)
+
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedArgs, result)
+		})
+	}
+}
+
+func TestParseFailure(t *testing.T) {
+	testCases := []struct {
+		desc          string
+		executable    *executable.Executable
+		environment   map[string]string
+		arguments     []string
+		expectedError string
+	}{
+		{
+			desc:          "It fails if SSH connection is not set",
+			executable:    &executable.Executable{Name: executable.GitlabShell},
+			arguments:     []string{},
+			expectedError: "Only SSH allowed",
+		},
+		{
+			desc:       "It fails if SSH command is invalid",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
+			},
+			arguments:     []string{},
+			expectedError: "Invalid SSH command",
+		},
+		{
+			desc:          "With not enough arguments for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user"},
+			expectedError: "# Insufficient arguments. 1. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>",
+		},
+		{
+			desc:          "With too many arguments for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "user", "key", "something-else"},
+			expectedError: "# Insufficient arguments. 4. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>",
+		},
+		{
+			desc:          "With missing username for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "", "key"},
+			expectedError: "# No username provided",
+		},
+		{
+			desc:          "With missing key for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "user", ""},
+			expectedError: "# No key provided",
+		},
+		{
+			desc:          "With not enough arguments for the AuthorizedPrincipalsCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedPrincipalsCheck},
+			arguments:     []string{"key"},
+			expectedError: "# Insufficient arguments. 1. Usage\n#\tgitlab-shell-authorized-principals-check <key-id> <principal1> [<principal2>...]",
+		},
+		{
+			desc:          "With missing key_id for the AuthorizedPrincipalsCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedPrincipalsCheck},
+			arguments:     []string{"", "principal"},
+			expectedError: "# No key_id provided",
+		},
+		{
+			desc:          "With blank principal for the AuthorizedPrincipalsCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedPrincipalsCheck},
+			arguments:     []string{"key", "principal", ""},
+			expectedError: "# An invalid principal was provided",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
+
+			_, err := Parse(tc.executable, tc.arguments)
+
+			require.EqualError(t, err, tc.expectedError)
+		})
+	}
+}
diff --git a/internal/command/commandargs/generic_args.go b/internal/command/commandargs/generic_args.go
new file mode 100644
index 0000000..96bed99
--- /dev/null
+++ b/internal/command/commandargs/generic_args.go
@@ -0,0 +1,14 @@
+package commandargs
+
+type GenericArgs struct {
+	Arguments []string
+}
+
+func (b *GenericArgs) Parse() error {
+	// Do nothing
+	return nil
+}
+
+func (b *GenericArgs) GetArguments() []string {
+	return b.Arguments
+}
diff --git a/internal/command/commandargs/shell.go b/internal/command/commandargs/shell.go
new file mode 100644
index 0000000..7e2b72e
--- /dev/null
+++ b/internal/command/commandargs/shell.go
@@ -0,0 +1,131 @@
+package commandargs
+
+import (
+	"errors"
+	"os"
+	"regexp"
+
+	"github.com/mattn/go-shellwords"
+)
+
+const (
+	Discover         CommandType = "discover"
+	TwoFactorRecover CommandType = "2fa_recovery_codes"
+	LfsAuthenticate  CommandType = "git-lfs-authenticate"
+	ReceivePack      CommandType = "git-receive-pack"
+	UploadPack       CommandType = "git-upload-pack"
+	UploadArchive    CommandType = "git-upload-archive"
+)
+
+var (
+	whoKeyRegex      = regexp.MustCompile(`\bkey-(?P<keyid>\d+)\b`)
+	whoUsernameRegex = regexp.MustCompile(`\busername-(?P<username>\S+)\b`)
+)
+
+type Shell struct {
+	Arguments      []string
+	GitlabUsername string
+	GitlabKeyId    string
+	SshArgs        []string
+	CommandType    CommandType
+}
+
+func (s *Shell) Parse() error {
+	if err := s.validate(); err != nil {
+		return err
+	}
+
+	s.parseWho()
+	s.defineCommandType()
+
+	return nil
+}
+
+func (s *Shell) GetArguments() []string {
+	return s.Arguments
+}
+
+func (s *Shell) validate() error {
+	if !s.isSshConnection() {
+		return errors.New("Only SSH allowed")
+	}
+
+	if !s.isValidSshCommand() {
+		return errors.New("Invalid SSH command")
+	}
+
+	return nil
+}
+
+func (s *Shell) isSshConnection() bool {
+	ok := os.Getenv("SSH_CONNECTION")
+	return ok != ""
+}
+
+func (s *Shell) isValidSshCommand() bool {
+	err := s.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
+	return err == nil
+}
+
+func (s *Shell) parseWho() {
+	for _, argument := range s.Arguments {
+		if keyId := tryParseKeyId(argument); keyId != "" {
+			s.GitlabKeyId = keyId
+			break
+		}
+
+		if username := tryParseUsername(argument); username != "" {
+			s.GitlabUsername = username
+			break
+		}
+	}
+}
+
+func tryParseKeyId(argument string) string {
+	matchInfo := whoKeyRegex.FindStringSubmatch(argument)
+	if len(matchInfo) == 2 {
+		// The first element is the full matched string
+		// The second element is the named `keyid`
+		return matchInfo[1]
+	}
+
+	return ""
+}
+
+func tryParseUsername(argument string) string {
+	matchInfo := whoUsernameRegex.FindStringSubmatch(argument)
+	if len(matchInfo) == 2 {
+		// The first element is the full matched string
+		// The second element is the named `username`
+		return matchInfo[1]
+	}
+
+	return ""
+}
+
+func (s *Shell) parseCommand(commandString string) error {
+	args, err := shellwords.Parse(commandString)
+	if err != nil {
+		return err
+	}
+
+	// Handle Git for Windows 2.14 using "git upload-pack" instead of git-upload-pack
+	if len(args) > 1 && args[0] == "git" {
+		command := args[0] + "-" + args[1]
+		commandArgs := args[2:]
+
+		args = append([]string{command}, commandArgs...)
+	}
+
+	s.SshArgs = args
+
+	return nil
+}
+
+func (s *Shell) defineCommandType() {
+	if len(s.SshArgs) == 0 {
+		s.CommandType = Discover
+	} else {
+		s.CommandType = CommandType(s.SshArgs[0])
+	}
+}