<commit_before>package jwthelper

import (
	"errors"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"github.com/northbright/errorhelper"
	"github.com/northbright/pathhelper"
	"io/ioutil"
	"net/http"
	"sync"
)

// Key struct consists of algorithm, signning key and verifying key.
type Key struct {
	Method    jwt.SigningMethod // jwt.SigningMethod
	SignKey   interface{}       // Signing key. HMAC: []byte, RSA / RSAPSS: *crypto/rsa.PrivateKey, ECDSA: *crypto/ecdsa.PrivateKey.
	VerifyKey interface{}       // Verifying key. HMAC: []byte, RSA / RSAPSS: *crypto/rsa.PublicKey, ECDSA: *crypto/ecdsa.PublicKey.
}

//  KeyManger manages the keys by using kid(key id).
type KeyManager struct {
	Keys         map[string]*Key // Key map. Key: kid(key id), Value: Key Struct
	sync.RWMutex                 // Access map concurrently.
}

const (
	AvailableAlgs string = "Available algs: HS256,HS384,HS512,RS256,RS384,RS512,PS256,PS384,PS512,ES256,ES384,ES512"
)

var (
	km KeyManager = KeyManager{Keys: make(map[string]*Key)} // Internal key manager.
)

// ReadKey() reads key bytes from the key file.
func ReadKey(keyFile string) (key []byte, err error) {
	if err := errorhelper.GenEmptyStringError(keyFile, "keyFile"); err != nil {
		return nil, err
	}

	// Make Abs key file path with current executable path if KeyFilePath is relative.
	p := ""
	if p, err = pathhelper.GetAbsPath(keyFile); err != nil {
		return nil, err
	}

	buf := []byte{}
	if buf, err = ioutil.ReadFile(p); err != nil {
		return nil, err
	}

	return buf, nil
}

// SetKey() sets the kid - Key pair.
//
//   Params:
//       kid: Key id. It should be unique.
//       key: Key struct.
func SetKey(kid string, key *Key) {
	km.Lock()
	km.Keys[kid] = key
	km.Unlock()
}

// GetKey() return the key struct by given kid.
func GetKey(kid string) (k *Key, err error) {
	if err := errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return nil, err
	}

	km.RLock()
	k, ok := km.Keys[kid]
	km.RUnlock()

	if !ok {
		return nil, errorhelper.GenMapKeyNotFoundError(kid)
	}

	return k, nil
}

// DeleteKey() deletes the specified entry from the key map.
func DeleteKey(kid string) (err error) {
	if _, err := GetKey(kid); err != nil {
		return err
	}

	km.Lock()
	delete(km.Keys, kid)
	km.Unlock()

	return nil
}

// SetKeyFromFile() reads the key files and stores the unique kid - Key information pair.
//
//   Params:
//       kid: Key id(unique).
//       alg: JWT alg.
//       signKeyFile: Signing key file.
//       verifyKeyFile: Verifying key file.
//   Return:
//       err: error.
//   Notes:
//       1. Current Available JWT "alg": HS256, HS384, HS512, RS256, RS384, RS512, PS256, PS384, PS512, ES256, ES384, ES512.
//       2. HMAC using SHA-XXX is a symmetric key algorithm. It just read signKeyFile as secret key(verifyKeyFile is ignored).
//       3. How to Generate Keys for JWT algs:
//          https://github.com/northbright/Notes/blob/master/jwt/generate_keys_for_jwt_alg.md
func SetKeyFromFile(kid, alg, signKeyFile, verifyKeyFile string) (err error) {
	key := &Key{}

	if err := errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return err
	}

	if err := errorhelper.GenEmptyStringError(alg, "alg"); err != nil {
		return err
	}

	m := jwt.GetSigningMethod(alg)
	if m == nil {
		msg := fmt.Sprintf("Incorrect alg: %s. %s", alg, AvailableAlgs)
		return errors.New(msg)
	}

	// Set Signing Method
	key.Method = m

	switch alg {
	case "HS256", "HS384", "HS512":
		if key.SignKey, err = ReadKey(signKeyFile); err != nil {
			return err
		}
		key.VerifyKey = key.SignKey

	case "RS256", "RS384", "RS512", "PS256", "PS384", "PS512":
		buf := []byte{}
		if buf, err = ReadKey(signKeyFile); err != nil {
			return err
		}

		if key.SignKey, err = jwt.ParseRSAPrivateKeyFromPEM(buf); err != nil {
			return err
		}

		if buf, err = ReadKey(verifyKeyFile); err != nil {
			return err
		}

		if key.VerifyKey, err = jwt.ParseRSAPublicKeyFromPEM(buf); err != nil {
			return err
		}

	case "ES256", "ES384", "ES512":
		buf := []byte{}
		if buf, err = ReadKey(signKeyFile); err != nil {
			return err
		}

		if key.SignKey, err = jwt.ParseECPrivateKeyFromPEM(buf); err != nil {
			return err
		}

		if buf, err = ReadKey(verifyKeyFile); err != nil {
			return err
		}

		if key.VerifyKey, err = jwt.ParseECPublicKeyFromPEM(buf); err != nil {
			return err
		}
	default:
		msg := fmt.Sprintf("Incorrect alg: %s. %s", alg, AvailableAlgs)
		return errors.New(msg)
	}

	SetKey(kid, key)
	return nil
}

// CreateTokenString() creates a new JWT token string.
//
//   Params:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//   Return:
//       tokenString: new created JWT token string.
//       err: error.
func CreateTokenString(kid string, claims map[string]interface{}) (tokenString string, err error) {
	var k *Key

	if err = errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return "", err
	}

	if k, err = GetKey(kid); err != nil {
		return "", err
	}

	t := jwt.New(k.Method)
	t.Header["kid"] = kid
	t.Claims = claims
	return t.SignedString(k.SignKey)
}

// jwt-go's KeyFunc type:
//
// type Keyfunc func(*Token) (interface{}, error)
func keyFunc(token *jwt.Token) (interface{}, error) {
	if _, ok := token.Header["kid"]; !ok {
		return nil, errorhelper.GenMapKeyNotFoundError("kid")
	}

	kid := ""
	if str, ok := token.Header["kid"].(string); !ok {
		msg := fmt.Sprintf("token.Header[\"kid\"]'s type is %T, but not string.", token.Header["kid"])
		return nil, errors.New(msg)
	} else {
		kid = str
	}

	key, err := GetKey(kid)
	if err != nil {
		return nil, err
	}

	// Check signing method
	if token.Method.Alg() != key.Method.Alg() {
		return nil, errors.New("Signing Method Error.")
	}

	return key.VerifyKey, nil
}

// Parse() parses and validates the input token string.
//
//   Params:
//       tokenString: input JWT token string.
//   Return:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//       valid: token is valid or not.
//       err: error.
func Parse(tokenString string) (kid string, claims map[string]interface{}, valid bool, err error) {
	t, err := jwt.Parse(tokenString, keyFunc)
	if err != nil {
		return "", nil, false, err
	}

	return t.Header["kid"].(string), t.Claims, t.Valid, nil
}

// ParseFromRequest() parses and validates the input token string in an http.Request. It's a wrapper of jwt.ParseFromRequest().
//
//   Params:
//       r: http.Request may contain jwt token.
//   Return:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//       valid: token is valid or not.
//       err: error.
func ParseFromRequest(r *http.Request) (kid string, claims map[string]interface{}, valid bool, err error) {
	t, err := jwt.ParseFromRequest(r, keyFunc)
	if err != nil {
		return "", nil, false, err
	}

	return t.Header["kid"].(string), t.Claims, t.Valid, nil
}
<commit_msg>Import non-stadard packages after standard ones.<commit_after>package jwthelper

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"

	"github.com/dgrijalva/jwt-go"
	"github.com/northbright/errorhelper"
	"github.com/northbright/pathhelper"
)

// Key struct consists of algorithm, signning key and verifying key.
type Key struct {
	Method    jwt.SigningMethod // jwt.SigningMethod
	SignKey   interface{}       // Signing key. HMAC: []byte, RSA / RSAPSS: *crypto/rsa.PrivateKey, ECDSA: *crypto/ecdsa.PrivateKey.
	VerifyKey interface{}       // Verifying key. HMAC: []byte, RSA / RSAPSS: *crypto/rsa.PublicKey, ECDSA: *crypto/ecdsa.PublicKey.
}

//  KeyManger manages the keys by using kid(key id).
type KeyManager struct {
	Keys         map[string]*Key // Key map. Key: kid(key id), Value: Key Struct
	sync.RWMutex                 // Access map concurrently.
}

const (
	AvailableAlgs string = "Available algs: HS256,HS384,HS512,RS256,RS384,RS512,PS256,PS384,PS512,ES256,ES384,ES512"
)

var (
	km KeyManager = KeyManager{Keys: make(map[string]*Key)} // Internal key manager.
)

// ReadKey() reads key bytes from the key file.
func ReadKey(keyFile string) (key []byte, err error) {
	if err := errorhelper.GenEmptyStringError(keyFile, "keyFile"); err != nil {
		return nil, err
	}

	// Make Abs key file path with current executable path if KeyFilePath is relative.
	p := ""
	if p, err = pathhelper.GetAbsPath(keyFile); err != nil {
		return nil, err
	}

	buf := []byte{}
	if buf, err = ioutil.ReadFile(p); err != nil {
		return nil, err
	}

	return buf, nil
}

// SetKey() sets the kid - Key pair.
//
//   Params:
//       kid: Key id. It should be unique.
//       key: Key struct.
func SetKey(kid string, key *Key) {
	km.Lock()
	km.Keys[kid] = key
	km.Unlock()
}

// GetKey() return the key struct by given kid.
func GetKey(kid string) (k *Key, err error) {
	if err := errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return nil, err
	}

	km.RLock()
	k, ok := km.Keys[kid]
	km.RUnlock()

	if !ok {
		return nil, errorhelper.GenMapKeyNotFoundError(kid)
	}

	return k, nil
}

// DeleteKey() deletes the specified entry from the key map.
func DeleteKey(kid string) (err error) {
	if _, err := GetKey(kid); err != nil {
		return err
	}

	km.Lock()
	delete(km.Keys, kid)
	km.Unlock()

	return nil
}

// SetKeyFromFile() reads the key files and stores the unique kid - Key information pair.
//
//   Params:
//       kid: Key id(unique).
//       alg: JWT alg.
//       signKeyFile: Signing key file.
//       verifyKeyFile: Verifying key file.
//   Return:
//       err: error.
//   Notes:
//       1. Current Available JWT "alg": HS256, HS384, HS512, RS256, RS384, RS512, PS256, PS384, PS512, ES256, ES384, ES512.
//       2. HMAC using SHA-XXX is a symmetric key algorithm. It just read signKeyFile as secret key(verifyKeyFile is ignored).
//       3. How to Generate Keys for JWT algs:
//          https://github.com/northbright/Notes/blob/master/jwt/generate_keys_for_jwt_alg.md
func SetKeyFromFile(kid, alg, signKeyFile, verifyKeyFile string) (err error) {
	key := &Key{}

	if err := errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return err
	}

	if err := errorhelper.GenEmptyStringError(alg, "alg"); err != nil {
		return err
	}

	m := jwt.GetSigningMethod(alg)
	if m == nil {
		msg := fmt.Sprintf("Incorrect alg: %s. %s", alg, AvailableAlgs)
		return errors.New(msg)
	}

	// Set Signing Method
	key.Method = m

	switch alg {
	case "HS256", "HS384", "HS512":
		if key.SignKey, err = ReadKey(signKeyFile); err != nil {
			return err
		}
		key.VerifyKey = key.SignKey

	case "RS256", "RS384", "RS512", "PS256", "PS384", "PS512":
		buf := []byte{}
		if buf, err = ReadKey(signKeyFile); err != nil {
			return err
		}

		if key.SignKey, err = jwt.ParseRSAPrivateKeyFromPEM(buf); err != nil {
			return err
		}

		if buf, err = ReadKey(verifyKeyFile); err != nil {
			return err
		}

		if key.VerifyKey, err = jwt.ParseRSAPublicKeyFromPEM(buf); err != nil {
			return err
		}

	case "ES256", "ES384", "ES512":
		buf := []byte{}
		if buf, err = ReadKey(signKeyFile); err != nil {
			return err
		}

		if key.SignKey, err = jwt.ParseECPrivateKeyFromPEM(buf); err != nil {
			return err
		}

		if buf, err = ReadKey(verifyKeyFile); err != nil {
			return err
		}

		if key.VerifyKey, err = jwt.ParseECPublicKeyFromPEM(buf); err != nil {
			return err
		}
	default:
		msg := fmt.Sprintf("Incorrect alg: %s. %s", alg, AvailableAlgs)
		return errors.New(msg)
	}

	SetKey(kid, key)
	return nil
}

// CreateTokenString() creates a new JWT token string.
//
//   Params:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//   Return:
//       tokenString: new created JWT token string.
//       err: error.
func CreateTokenString(kid string, claims map[string]interface{}) (tokenString string, err error) {
	var k *Key

	if err = errorhelper.GenEmptyStringError(kid, "kid"); err != nil {
		return "", err
	}

	if k, err = GetKey(kid); err != nil {
		return "", err
	}

	t := jwt.New(k.Method)
	t.Header["kid"] = kid
	t.Claims = claims
	return t.SignedString(k.SignKey)
}

// jwt-go's KeyFunc type:
//
// type Keyfunc func(*Token) (interface{}, error)
func keyFunc(token *jwt.Token) (interface{}, error) {
	if _, ok := token.Header["kid"]; !ok {
		return nil, errorhelper.GenMapKeyNotFoundError("kid")
	}

	kid := ""
	if str, ok := token.Header["kid"].(string); !ok {
		msg := fmt.Sprintf("token.Header[\"kid\"]'s type is %T, but not string.", token.Header["kid"])
		return nil, errors.New(msg)
	} else {
		kid = str
	}

	key, err := GetKey(kid)
	if err != nil {
		return nil, err
	}

	// Check signing method
	if token.Method.Alg() != key.Method.Alg() {
		return nil, errors.New("Signing Method Error.")
	}

	return key.VerifyKey, nil
}

// Parse() parses and validates the input token string.
//
//   Params:
//       tokenString: input JWT token string.
//   Return:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//       valid: token is valid or not.
//       err: error.
func Parse(tokenString string) (kid string, claims map[string]interface{}, valid bool, err error) {
	t, err := jwt.Parse(tokenString, keyFunc)
	if err != nil {
		return "", nil, false, err
	}

	return t.Header["kid"].(string), t.Claims, t.Valid, nil
}

// ParseFromRequest() parses and validates the input token string in an http.Request. It's a wrapper of jwt.ParseFromRequest().
//
//   Params:
//       r: http.Request may contain jwt token.
//   Return:
//       kid: Key id.
//       claims: map[string]interface{} to fill the jwt.Token[Claims].
//       valid: token is valid or not.
//       err: error.
func ParseFromRequest(r *http.Request) (kid string, claims map[string]interface{}, valid bool, err error) {
	t, err := jwt.ParseFromRequest(r, keyFunc)
	if err != nil {
		return "", nil, false, err
	}

	return t.Header["kid"].(string), t.Claims, t.Valid, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package main

import (
	"github.com/limetext/lime/backend"
	"github.com/limetext/termbox-go"
	. "github.com/limetext/text"
	"testing"
	"time"
)

func TestPadLineRunes(t *testing.T) {
	var testPadData = []struct {
		line     int
		total    int
		expected []rune
	}{
		{1, 3, []rune{' ', ' ', '1', ' '}},
		{10, 3, []rune{' ', '1', '0', ' '}},
		{100, 3, []rune{'1', '0', '0', ' '}},
	}

	for _, p := range testPadData {
		runes := intToRunes(p.line)
		padded := padLineRunes(runes, p.total)

		for i, r := range p.expected {
			if r != padded[i] {
				t.Error("Expected runes to be padded")
			}
		}
	}
}

func TestIntToRunes(t *testing.T) {
	actual := intToRunes(100)
	expected := []rune{'1', '0', '0'}
	for i, r := range expected {
		if actual[i] != r {
			t.Error("Expected int to be transformed to runes")
		}
	}
}

func TestGetCaretStyle(t *testing.T) {
	var testcases = []struct {
		style    string
		inverse  bool
		expected termbox.Attribute
	}{
		{"", false, termbox.AttrUnderline},
		{"block", false, termbox.AttrReverse},
		{"", true, termbox.AttrReverse},
	}

	for _, tc := range testcases {
		if style := getCaretStyle(tc.style, tc.inverse); style != tc.expected {
			t.Errorf("Expected %s, got %s", tc.expected, style)
		}
	}
}

func TestUpdateVisibleRegion(t *testing.T) {
	var (
		fe tbfe
		e  = backend.GetEditor()
		w  = e.NewWindow()
		v  = w.NewFile()
	)

	fe.layout = make(map[*backend.View]layout)
	fe.layout[v] = layout{0, 0, 100, 100 - *consoleHeight - 1, Region{}, 0}
	fe.setupCallbacks(v)

	edit := v.BeginEdit()
	v.Insert(edit, 0, "foo")
	v.EndEdit(edit)

	if end := fe.layout[v].visible.End(); end != 3 {
		t.Fatalf("Expected 3, got %d", end)
	}
}

func TestCreateFrontend(t *testing.T) {
	var frontend *tbfe
	frontendWasCreated := make(chan bool, 0)
	go func() {
		frontend = createFrontend()
		frontendWasCreated <- true
	}()

	select {
	case <-frontendWasCreated:
		break
	case <-time.After(2 * time.Second):
		t.Error("Frontend was not created within timeout")
	}

	frontend.lock.Lock()
	*showConsole = true
	if frontend.editor == nil {
		t.Error("Editor is nil")
	}

	if frontend.console == nil {
		t.Error("Current console is nil")
	}

	if frontend.currentWindow == nil {
		t.Error("Current window is nil")
	}

	if frontend.currentView == nil {
		t.Error("Current view is nil")
	}

	if _, ok := frontend.layout[frontend.currentView]; !ok {
		t.Error("Current view not in layout")
	}

	if _, ok := frontend.layout[frontend.console]; !ok {
		t.Error("Console view not in layout")
	}

	if len(frontend.layout) != 2 {
		t.Error("Layout too big")
	}
	frontend.lock.Unlock()

}

func TestLoopShutdown(t *testing.T) {
	frontend := createFrontend()

	loopHasExited := make(chan bool, 0)
	go func() {
		frontend.loop()
		loopHasExited <- true
	}()

	frontend.shutdown <- true

	select {
	case <-loopHasExited:
		break
	case <-time.After(2 * time.Second):
		t.Error("Loop did not terminate within timeout")
	}
}
<commit_msg>[frontend/termbox] Add some extra information to a failed test.<commit_after>// Copyright 2013 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package main

import (
	"github.com/limetext/lime/backend"
	"github.com/limetext/termbox-go"
	. "github.com/limetext/text"
	"testing"
	"time"
)

func TestPadLineRunes(t *testing.T) {
	var testPadData = []struct {
		line     int
		total    int
		expected []rune
	}{
		{1, 3, []rune{' ', ' ', '1', ' '}},
		{10, 3, []rune{' ', '1', '0', ' '}},
		{100, 3, []rune{'1', '0', '0', ' '}},
	}

	for _, p := range testPadData {
		runes := intToRunes(p.line)
		padded := padLineRunes(runes, p.total)

		for i, r := range p.expected {
			if r != padded[i] {
				t.Error("Expected runes to be padded")
			}
		}
	}
}

func TestIntToRunes(t *testing.T) {
	actual := intToRunes(100)
	expected := []rune{'1', '0', '0'}
	for i, r := range expected {
		if actual[i] != r {
			t.Error("Expected int to be transformed to runes")
		}
	}
}

func TestGetCaretStyle(t *testing.T) {
	var testcases = []struct {
		style    string
		inverse  bool
		expected termbox.Attribute
	}{
		{"", false, termbox.AttrUnderline},
		{"block", false, termbox.AttrReverse},
		{"", true, termbox.AttrReverse},
	}

	for _, tc := range testcases {
		if style := getCaretStyle(tc.style, tc.inverse); style != tc.expected {
			t.Errorf("Expected %s, got %s", tc.expected, style)
		}
	}
}

func TestUpdateVisibleRegion(t *testing.T) {
	var (
		fe tbfe
		e  = backend.GetEditor()
		w  = e.NewWindow()
		v  = w.NewFile()
	)

	fe.layout = make(map[*backend.View]layout)
	fe.layout[v] = layout{0, 0, 100, 100 - *consoleHeight - 1, Region{}, 0}
	fe.setupCallbacks(v)

	edit := v.BeginEdit()
	v.Insert(edit, 0, "foo")
	v.EndEdit(edit)

	if end := fe.layout[v].visible.End(); end != 3 {
		t.Fatalf("Expected 3, got %d", end)
	}
}

func TestCreateFrontend(t *testing.T) {
	var frontend *tbfe
	frontendWasCreated := make(chan bool, 0)
	go func() {
		frontend = createFrontend()
		frontendWasCreated <- true
	}()

	select {
	case <-frontendWasCreated:
		break
	case <-time.After(2 * time.Second):
		t.Error("Frontend was not created within timeout")
	}

	frontend.lock.Lock()
	defer frontend.lock.Unlock()

	*showConsole = true
	if frontend.editor == nil {
		t.Error("Editor is nil")
	}

	if frontend.console == nil {
		t.Error("Current console is nil")
	}

	if frontend.currentWindow == nil {
		t.Error("Current window is nil")
	}

	if frontend.currentView == nil {
		t.Error("Current view is nil")
	}

	if _, ok := frontend.layout[frontend.currentView]; !ok {
		t.Error("Current view not in layout")
	}

	if _, ok := frontend.layout[frontend.console]; !ok {
		t.Error("Console view not in layout")
	}

	if len(frontend.layout) != 2 {
		t.Errorf("Layout has length %d, but should have length 2", len(frontend.layout))
	}
}

func TestLoopShutdown(t *testing.T) {
	frontend := createFrontend()

	loopHasExited := make(chan bool, 0)
	go func() {
		frontend.loop()
		loopHasExited <- true
	}()

	frontend.shutdown <- true

	select {
	case <-loopHasExited:
		break
	case <-time.After(2 * time.Second):
		t.Error("Loop did not terminate within timeout")
	}
}
<|endoftext|>


<commit_before>package creds

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const (
	mfaCodeRegexString = `^\d{6}$`
)

var mfaCodeRegex = regexp.MustCompile(mfaCodeRegexString)

// MfaPrompt defines an object which recieves an Mfa ARN and returns an Mfa code
type MfaPrompt interface {
	Prompt(string) (string, error)
}

// revive:disable-next-line:flag-parameter
func (c Creds) handleMfa(useMfa bool, mfaCode string, mfaPrompt MfaPrompt) (string, string, error) {
	logger.InfoMsg("handling mfa options")

	if !useMfa && mfaCode == "" {
		logger.InfoMsg("mfa is disabled")
		return "", "", nil
	}
	mfaSerial, err := c.MfaArn()
	if err != nil {
		return "", "", err
	}
	if mfaCode != "" {
		logger.InfoMsg("mfa code already provided")
		err := validateMfa(mfaCode)
		if err != nil {
			return "", "", err
		}

		return mfaCode, mfaSerial, nil
	}
	if mfaPrompt == nil {
		logger.InfoMsg("using default mfa prompt")
		// revive:disable-next-line:modifies-parameter
		mfaPrompt = &DefaultMfaPrompt{}
	}
	logger.InfoMsg("prompting for mfa code")
	// revive:disable-next-line:modifies-parameter
	mfaCode, err = mfaPrompt.Prompt(mfaSerial)
	if err != nil {
		return "", "", err
	}
	return mfaCode, mfaSerial, nil
}

// DefaultMfaPrompt defines the standard CLI-based MFA prompt
type DefaultMfaPrompt struct{}

// Prompt asks the user for their MFA token
func (p *DefaultMfaPrompt) Prompt(_ string) (string, error) {
	mfaReader := bufio.NewReader(os.Stdin)
	fmt.Fprint(os.Stderr, "MFA Code: ")
	mfa, err := mfaReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	trimmed := strings.TrimSpace(mfa)
	err = validateMfa(trimmed)
	if err != nil {
		return "", err
	}
	return trimmed, nil
}

func validateMfa(code string) error {
	logger.InfoMsg("validating mfa")
	if mfaCodeRegex.MatchString(code) {
		return nil
	}
	return fmt.Errorf("provided mfa code does not match the necessary format")
}

// MultiMfaPrompt allows a slice of sequential backends to check for Mfa
type MultiMfaPrompt struct {
	Backends []MfaPrompt
}

// Prompt iterates through the backends to find an Mfa code
func (m *MultiMfaPrompt) Prompt(arn string) (string, error) {
	logger.InfoMsgf("looking up %s in multi mfa prompt", arn)
	for index, item := range m.Backends {
		res, err := item.Prompt(arn)
		if err == nil {
			return res, nil
		}
		logger.InfoMsgf("failed backend lookup %d with %s", index, err)
	}
	return "", fmt.Errorf("all backends failed to return mfa code")
}
<commit_msg>add writable MFA prompt<commit_after>package creds

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const (
	mfaCodeRegexString = `^\d{6}$`
)

var mfaCodeRegex = regexp.MustCompile(mfaCodeRegexString)

// MfaPrompt defines an object which recieves an Mfa ARN and returns an Mfa code
type MfaPrompt interface {
	Prompt(string) (string, error)
}

// WritableMfaPrompt defines an MFA Prompt which can store a new secret
type WritableMfaPrompt interface {
	Store(string, string) error
	RetryText(string) string
}

// revive:disable-next-line:flag-parameter
func (c Creds) handleMfa(useMfa bool, mfaCode string, mfaPrompt MfaPrompt) (string, string, error) {
	logger.InfoMsg("handling mfa options")

	if !useMfa && mfaCode == "" {
		logger.InfoMsg("mfa is disabled")
		return "", "", nil
	}
	mfaSerial, err := c.MfaArn()
	if err != nil {
		return "", "", err
	}
	if mfaCode != "" {
		logger.InfoMsg("mfa code already provided")
		err := validateMfa(mfaCode)
		if err != nil {
			return "", "", err
		}

		return mfaCode, mfaSerial, nil
	}
	if mfaPrompt == nil {
		logger.InfoMsg("using default mfa prompt")
		// revive:disable-next-line:modifies-parameter
		mfaPrompt = &DefaultMfaPrompt{}
	}
	logger.InfoMsg("prompting for mfa code")
	// revive:disable-next-line:modifies-parameter
	mfaCode, err = mfaPrompt.Prompt(mfaSerial)
	if err != nil {
		return "", "", err
	}
	return mfaCode, mfaSerial, nil
}

// DefaultMfaPrompt defines the standard CLI-based MFA prompt
type DefaultMfaPrompt struct{}

// Prompt asks the user for their MFA token
func (p *DefaultMfaPrompt) Prompt(_ string) (string, error) {
	mfaReader := bufio.NewReader(os.Stdin)
	fmt.Fprint(os.Stderr, "MFA Code: ")
	mfa, err := mfaReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	trimmed := strings.TrimSpace(mfa)
	err = validateMfa(trimmed)
	if err != nil {
		return "", err
	}
	return trimmed, nil
}

func validateMfa(code string) error {
	logger.InfoMsg("validating mfa")
	if mfaCodeRegex.MatchString(code) {
		return nil
	}
	return fmt.Errorf("provided mfa code does not match the necessary format")
}

// MultiMfaPrompt allows a slice of sequential backends to check for Mfa
type MultiMfaPrompt struct {
	Backends []MfaPrompt
}

// Prompt iterates through the backends to find an Mfa code
func (m *MultiMfaPrompt) Prompt(arn string) (string, error) {
	logger.InfoMsgf("looking up %s in multi mfa prompt", arn)
	for index, item := range m.Backends {
		res, err := item.Prompt(arn)
		if err == nil {
			return res, nil
		}
		logger.InfoMsgf("failed backend lookup %d with %s", index, err)
	}
	return "", fmt.Errorf("all backends failed to return mfa code")
}

// Store attempts to store the Mfa in a backend
func (m *MultiMfaPrompt) Store(arn, seed string) error {
	logger.InfoMsgf("attempting to store %s in multi mfa object", arn)
	writer := m.getWriter()
	if writer == nil {
		return fmt.Errorf("no writable mfa backends found")
	}
	return writer.Store(arn, seed)
}

// RetryText returns helper text for retrying a failed mfa storage
func (m *MultiMfaPrompt) RetryText(arn string) string {
	writer := m.getWriter()
	if writer == nil {
		return ""
	}
	return writer.RetryText(arn)
}

func (m *MultiMfaPrompt) getWriter() *creds.WritableMfaPrompt {
	for item := range m.Backends {
		writer, ok := item.(WritableMfaPrompt)
		if ok {
			logger.InfoMsgf("found mfa writer with index %d", index)
			return writer
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"code.google.com/p/go.crypto/ssh"

	"code.google.com/p/gopass"
	"github.com/codegangsta/cli"
)

const ENV_PATTERN string = "^[A-Za-z_][A-Za-z0-9_]*=.*"

func decryptPEM(pemblock *pem.Block, filename string) ([]byte, error) {
	var err error
	if _, err = fmt.Fprintf(os.Stderr, "Enter passphrase for %s: ", filename); err != nil {
		return []byte(""), err
	}

	// we already emit the prompt to stderr; GetPass only emits to stdout
	var passwd string
	passwd, err = gopass.GetPass("")
	fmt.Fprintln(os.Stderr, "")
	if err != nil {
		return []byte(""), err
	}

	var decryptedBytes []byte
	if decryptedBytes, err = x509.DecryptPEMBlock(pemblock, []byte(passwd)); err != nil {
		return []byte(""), err
	}

	pemBytes := pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: decryptedBytes,
	}
	decryptedPEM := pem.EncodeToMemory(&pemBytes)
	return decryptedPEM, nil
}

func getPrivateKey(c *cli.Context) (filename string) {
	if c.String("key") == "" {
		filename = filepath.Join(os.Getenv("HOME"), "/.ssh/id_rsa")
	} else {
		filename = c.String("key")
	}
	return filename
}

func splitUserAndAccount(arg string) (string, string, error) {
	atpos := strings.LastIndex(arg, "@")
	if atpos < 1 {
		err := errors.New("Invalid account format; please specify <username>@<account>")
		return "", "", err
	}
	// pull off everything before the last '@'
	return arg[atpos+1:], arg[0:atpos], nil
}

func getAccountAndUserName(c *cli.Context) (string, string, error) {
	if len(c.Args()) > 0 {
		user, acct, err := splitUserAndAccount(c.Args()[0])
		if err != nil {
			return "", "", err
		}
		return user, acct, nil
	}
	if c.String("credentials") != "" {
		user, acct, err := splitUserAndAccount(c.String("credentials"))
		if err != nil {
			return "", "", err
		}
		return user, acct, nil
	} else {
		return c.String("account"), c.String("username"), nil
	}
}

func parseUserAndAccount(c *cli.Context) (username string, account string, err error) {
	if (c.String("username") == "" || c.String("account") == "") && c.Bool("force") {
		err = errors.New("Must specify both username and account with force")
		return "", "", err
	}

	// if username OR account were specified, but not both, complain
	if (c.String("username") != "" && c.String("account") == "") ||
		(c.String("username") == "" && c.String("account") != "") {
		if c.Bool("force") {
			err = errors.New("Must specify both username and account for force save")
		} else {
			err = errors.New("Must use force save when specifying username or account")
		}
		return "", "", err
	}

	// if username/account were specified, but force wasn't set, complain
	if c.String("username") != "" && c.String("account") != "" {
		if !c.Bool("force") {
			err = errors.New("Cannot specify username and/or account without force")
			return "", "", err
		} else {
			log.Print("WARNING: saving credentials without verifying username or account alias")
			username = c.String("username")
			account = c.String("account")
		}
	}
	return username, account, nil
}

func parseEnvironmentArgs(c *cli.Context) (map[string]string, error) {
	if len(c.StringSlice("env")) == 0 {
		return nil, nil
	}

	envMap := make(map[string]string)
	for _, arg := range c.StringSlice("env") {
		match, err := regexp.Match(ENV_PATTERN, []byte(arg))
		if err != nil {
			return nil, err
		}
		if !match {
			log.Print("WARNING: Skipping env argument " + arg + " -- not in NAME=value format")
			continue
		}
		parts := strings.SplitN(arg, "=", 2)
		envMap[parts[0]] = parts[1]
	}
	return envMap, nil
}

func readSSHPubkeyFile(filename string) (pubkey ssh.PublicKey, err error) {
	pubkeyString, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	pubkey, _, _, _, err = ssh.ParseAuthorizedKey([]byte(pubkeyString))
	if err != nil {
		return nil, err
	}
	return pubkey, nil
}

func parseKeyArgs(c *cli.Context) (pubkeys []ssh.PublicKey, err error) {
	// no args, so just use the default
	if len(c.StringSlice("key")) == 0 {
		pubkey, err := readSSHPubkeyFile(filepath.Join(os.Getenv("HOME"), "/.ssh/id_rsa.pub"))
		if err != nil {
			return nil, err
		}
		pubkeys = append(pubkeys, pubkey)
		return pubkeys, nil
	}

	for _, arg := range c.StringSlice("key") {
		pubkey, err := readSSHPubkeyFile(arg)
		if err != nil {
			return nil, err
		}
		pubkeys = append(pubkeys, pubkey)
	}
	return pubkeys, nil
}

// parseLifetimeArgs attempts to be a little clever in determining what credential
// lifetime you've chosen. It returns a number of hours and an error. It assumes that
// the argument was passed in as hours.
func parseLifetimeArgs(c *cli.Context) (lifetime int, err error) {
	// the default is zero, which is our default
	if c.Int("lifetime") < 0 {
		return 0, nil
	}

	return c.Int("lifetime"), nil
}

func parseSaveArgs(c *cli.Context) (cred Credential, username, account string, pubkeys []ssh.PublicKey, lifetime int, err error) {
	pubkeys, err = parseKeyArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	username, account, err = parseUserAndAccount(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	envmap, err := parseEnvironmentArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	lifetime, err = parseLifetimeArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
	AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if AWSAccessKeyId == "" || AWSSecretAccessKey == "" {
		err := errors.New("Can't save, no credentials in the environment")
		if err != nil {
			return Credential{}, "", "", nil, 0, err
		}
	}
	cred = Credential{
		KeyId:     AWSAccessKeyId,
		SecretKey: AWSSecretAccessKey,
		EnvVars:   envmap,
	}

	return cred, username, account, pubkeys, lifetime, nil
}

func main() {
	app := cli.NewApp()
	app.Name = "credulous"
	app.Usage = "Secure AWS Credential Management"
	app.Version = "0.2.2"

	app.Commands = []cli.Command{
		{
			Name:  "save",
			Usage: "Save AWS credentials",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "key, k",
					Value: &cli.StringSlice{},
					Usage: "\n        SSH public keys",
				},
				cli.StringSliceFlag{
					Name:  "env, e",
					Value: &cli.StringSlice{},
					Usage: "\n        Environment variables to set in the form VAR=value",
				},
				cli.IntFlag{
					Name:  "lifetime, l",
					Value: 0,
					Usage: "\n        Credential lifetime in seconds (0 means forever)",
				},
				cli.BoolFlag{
					Name: "force, f",
					Usage: "\n        Force saving without validating username or account." +
						"\n        You MUST specify -u username -a account",
				},
				cli.StringFlag{
					Name:  "username, u",
					Value: "",
					Usage: "\n        Username (for use with '--force')",
				},
				cli.StringFlag{
					Name:  "account, a",
					Value: "",
					Usage: "\n        Account alias (for use with '--force')",
				},
			},
			Action: func(c *cli.Context) {
				cred, username, account, pubkeys, lifetime, err := parseSaveArgs(c)
				panic_the_err(err)
				err = SaveCredentials(cred, username, account, pubkeys, lifetime, c.Bool("force"))
				panic_the_err(err)
			},
		},

		{
			Name:  "source",
			Usage: "Source AWS credentials",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "account, a",
					Value: "",
					Usage: "\n        AWS Account alias or id",
				},
				cli.StringFlag{
					Name:  "key, k",
					Value: "",
					Usage: "\n        SSH private key",
				},
				cli.StringFlag{
					Name:  "username, u",
					Value: "",
					Usage: "\n        IAM User",
				},
				cli.StringFlag{
					Name:  "credentials, c",
					Value: "",
					Usage: "\n        Credentials, for example username@account",
				},
				cli.BoolFlag{
					Name:  "force, f",
					Usage: "\n        Force sourcing of credentials without validating username or account",
				},
			},
			Action: func(c *cli.Context) {
				keyfile := getPrivateKey(c)
				account, username, err := getAccountAndUserName(c)
				if err != nil {
					panic_the_err(err)
				}
				cred, err := RetrieveCredentials(account, username, keyfile)
				if err != nil {
					panic_the_err(err)
				}

				if !c.Bool("force") {
					err = cred.ValidateCredentials(account, username)
					if err != nil {
						panic_the_err(err)
					}
				}
				cred.Display(os.Stdout)
			},
		},

		{
			Name:  "current",
			Usage: "Show the username and alias of the currently-loaded credentials",
			Action: func(c *cli.Context) {
				AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
				AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
				if AWSAccessKeyId == "" || AWSSecretAccessKey == "" {
					err := errors.New("No amazon credentials are currently in your environment")
					panic_the_err(err)
				}
				cred := Credential{
					KeyId:     AWSAccessKeyId,
					SecretKey: AWSSecretAccessKey,
				}
				username, alias, err := getAWSUsernameAndAlias(cred)
				if err != nil {
					panic_the_err(err)
				}
				fmt.Printf("%s@%s\n", username, alias)
			},
		},

		{
			Name:  "display",
			Usage: "Display loaded AWS credentials",
			Action: func(c *cli.Context) {
				AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
				AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
				fmt.Printf("AWS_ACCESS_KEY_ID: %s\n", AWSAccessKeyId)
				fmt.Printf("AWS_SECRET_ACCESS_KEY: %s\n", AWSSecretAccessKey)
			},
		},

		{
			Name:  "list",
			Usage: "List available AWS credentials",
			Action: func(c *cli.Context) {
				rootDir, err := os.Open(getRootPath())
				if err != nil {
					panic_the_err(err)
				}
				set, err := listAvailableCredentials(rootDir)
				if err != nil {
					panic_the_err(err)
				}
				for _, cred := range set {
					fmt.Println(cred)
				}
			},
		},

		{
			Name:  "rotate",
			Usage: "Rotate current AWS credentials, deleting the oldest",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "lifetime, l",
					Value: 0,
					Usage: "\n        New credential lifetime in seconds (0 means forever)",
				},
				cli.StringFlag{
					Name:  "key, k",
					Value: "",
					Usage: "\n        SSH private key",
				},
				cli.StringSliceFlag{
					Name:  "env, e",
					Value: &cli.StringSlice{},
					Usage: "\n        Environment variables to set in the form VAR=value",
				},
			},
			Action: func(c *cli.Context) {
				cred, _, _, pubkey, lifetime, err := parseSaveArgs(c)
				panic_the_err(err)
				username, account, err := getAWSUsernameAndAlias(cred)
				panic_the_err(err)
				err = (&cred).rotateCredentials(username)
				panic_the_err(err)
				username, account, err = getAWSUsernameAndAlias(cred)
				panic_the_err(err)
				err = SaveCredentials(cred, username, account, pubkey, lifetime, c.Bool("force"))
				panic_the_err(err)
			},
		},
	}

	app.Run(os.Args)
}
<commit_msg>Support multiple pubkeys for encryption<commit_after>package main

import (
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"code.google.com/p/go.crypto/ssh"

	"code.google.com/p/gopass"
	"github.com/codegangsta/cli"
)

const ENV_PATTERN string = "^[A-Za-z_][A-Za-z0-9_]*=.*"

func decryptPEM(pemblock *pem.Block, filename string) ([]byte, error) {
	var err error
	if _, err = fmt.Fprintf(os.Stderr, "Enter passphrase for %s: ", filename); err != nil {
		return []byte(""), err
	}

	// we already emit the prompt to stderr; GetPass only emits to stdout
	var passwd string
	passwd, err = gopass.GetPass("")
	fmt.Fprintln(os.Stderr, "")
	if err != nil {
		return []byte(""), err
	}

	var decryptedBytes []byte
	if decryptedBytes, err = x509.DecryptPEMBlock(pemblock, []byte(passwd)); err != nil {
		return []byte(""), err
	}

	pemBytes := pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: decryptedBytes,
	}
	decryptedPEM := pem.EncodeToMemory(&pemBytes)
	return decryptedPEM, nil
}

func getPrivateKey(c *cli.Context) (filename string) {
	if c.String("key") == "" {
		filename = filepath.Join(os.Getenv("HOME"), "/.ssh/id_rsa")
	} else {
		filename = c.String("key")
	}
	return filename
}

func splitUserAndAccount(arg string) (string, string, error) {
	atpos := strings.LastIndex(arg, "@")
	if atpos < 1 {
		err := errors.New("Invalid account format; please specify <username>@<account>")
		return "", "", err
	}
	// pull off everything before the last '@'
	return arg[atpos+1:], arg[0:atpos], nil
}

func getAccountAndUserName(c *cli.Context) (string, string, error) {
	if len(c.Args()) > 0 {
		user, acct, err := splitUserAndAccount(c.Args()[0])
		if err != nil {
			return "", "", err
		}
		return user, acct, nil
	}
	if c.String("credentials") != "" {
		user, acct, err := splitUserAndAccount(c.String("credentials"))
		if err != nil {
			return "", "", err
		}
		return user, acct, nil
	} else {
		return c.String("account"), c.String("username"), nil
	}
}

func parseUserAndAccount(c *cli.Context) (username string, account string, err error) {
	if (c.String("username") == "" || c.String("account") == "") && c.Bool("force") {
		err = errors.New("Must specify both username and account with force")
		return "", "", err
	}

	// if username OR account were specified, but not both, complain
	if (c.String("username") != "" && c.String("account") == "") ||
		(c.String("username") == "" && c.String("account") != "") {
		if c.Bool("force") {
			err = errors.New("Must specify both username and account for force save")
		} else {
			err = errors.New("Must use force save when specifying username or account")
		}
		return "", "", err
	}

	// if username/account were specified, but force wasn't set, complain
	if c.String("username") != "" && c.String("account") != "" {
		if !c.Bool("force") {
			err = errors.New("Cannot specify username and/or account without force")
			return "", "", err
		} else {
			log.Print("WARNING: saving credentials without verifying username or account alias")
			username = c.String("username")
			account = c.String("account")
		}
	}
	return username, account, nil
}

func parseEnvironmentArgs(c *cli.Context) (map[string]string, error) {
	if len(c.StringSlice("env")) == 0 {
		return nil, nil
	}

	envMap := make(map[string]string)
	for _, arg := range c.StringSlice("env") {
		match, err := regexp.Match(ENV_PATTERN, []byte(arg))
		if err != nil {
			return nil, err
		}
		if !match {
			log.Print("WARNING: Skipping env argument " + arg + " -- not in NAME=value format")
			continue
		}
		parts := strings.SplitN(arg, "=", 2)
		envMap[parts[0]] = parts[1]
	}
	return envMap, nil
}

func readSSHPubkeyFile(filename string) (pubkey ssh.PublicKey, err error) {
	pubkeyString, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	pubkey, _, _, _, err = ssh.ParseAuthorizedKey([]byte(pubkeyString))
	if err != nil {
		return nil, err
	}
	return pubkey, nil
}

func parseKeyArgs(c *cli.Context) (pubkeys []ssh.PublicKey, err error) {
	// no args, so just use the default
	if len(c.StringSlice("key")) == 0 {
		pubkey, err := readSSHPubkeyFile(filepath.Join(os.Getenv("HOME"), "/.ssh/id_rsa.pub"))
		if err != nil {
			return nil, err
		}
		pubkeys = append(pubkeys, pubkey)
		return pubkeys, nil
	}

	for _, arg := range c.StringSlice("key") {
		pubkey, err := readSSHPubkeyFile(arg)
		if err != nil {
			return nil, err
		}
		pubkeys = append(pubkeys, pubkey)
	}
	return pubkeys, nil
}

// parseLifetimeArgs attempts to be a little clever in determining what credential
// lifetime you've chosen. It returns a number of hours and an error. It assumes that
// the argument was passed in as hours.
func parseLifetimeArgs(c *cli.Context) (lifetime int, err error) {
	// the default is zero, which is our default
	if c.Int("lifetime") < 0 {
		return 0, nil
	}

	return c.Int("lifetime"), nil
}

func parseSaveArgs(c *cli.Context) (cred Credential, username, account string, pubkeys []ssh.PublicKey, lifetime int, err error) {
	pubkeys, err = parseKeyArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	username, account, err = parseUserAndAccount(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	envmap, err := parseEnvironmentArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	lifetime, err = parseLifetimeArgs(c)
	if err != nil {
		return Credential{}, "", "", nil, 0, err
	}

	AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
	AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if AWSAccessKeyId == "" || AWSSecretAccessKey == "" {
		err := errors.New("Can't save, no credentials in the environment")
		if err != nil {
			return Credential{}, "", "", nil, 0, err
		}
	}
	cred = Credential{
		KeyId:     AWSAccessKeyId,
		SecretKey: AWSSecretAccessKey,
		EnvVars:   envmap,
	}

	return cred, username, account, pubkeys, lifetime, nil
}

func main() {
	app := cli.NewApp()
	app.Name = "credulous"
	app.Usage = "Secure AWS Credential Management"
	app.Version = "0.2.2"

	app.Commands = []cli.Command{
		{
			Name:  "save",
			Usage: "Save AWS credentials",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "key, k",
					Value: &cli.StringSlice{},
					Usage: "\n        SSH public keys for encryption",
				},
				cli.StringSliceFlag{
					Name:  "env, e",
					Value: &cli.StringSlice{},
					Usage: "\n        Environment variables to set in the form VAR=value",
				},
				cli.IntFlag{
					Name:  "lifetime, l",
					Value: 0,
					Usage: "\n        Credential lifetime in seconds (0 means forever)",
				},
				cli.BoolFlag{
					Name: "force, f",
					Usage: "\n        Force saving without validating username or account." +
						"\n        You MUST specify -u username -a account",
				},
				cli.StringFlag{
					Name:  "username, u",
					Value: "",
					Usage: "\n        Username (for use with '--force')",
				},
				cli.StringFlag{
					Name:  "account, a",
					Value: "",
					Usage: "\n        Account alias (for use with '--force')",
				},
			},
			Action: func(c *cli.Context) {
				cred, username, account, pubkeys, lifetime, err := parseSaveArgs(c)
				panic_the_err(err)
				err = SaveCredentials(cred, username, account, pubkeys, lifetime, c.Bool("force"))
				panic_the_err(err)
			},
		},

		{
			Name:  "source",
			Usage: "Source AWS credentials",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "account, a",
					Value: "",
					Usage: "\n        AWS Account alias or id",
				},
				cli.StringFlag{
					Name:  "key, k",
					Value: "",
					Usage: "\n        SSH private key",
				},
				cli.StringFlag{
					Name:  "username, u",
					Value: "",
					Usage: "\n        IAM User",
				},
				cli.StringFlag{
					Name:  "credentials, c",
					Value: "",
					Usage: "\n        Credentials, for example username@account",
				},
				cli.BoolFlag{
					Name:  "force, f",
					Usage: "\n        Force sourcing of credentials without validating username or account",
				},
			},
			Action: func(c *cli.Context) {
				keyfile := getPrivateKey(c)
				account, username, err := getAccountAndUserName(c)
				if err != nil {
					panic_the_err(err)
				}
				creds, err := RetrieveCredentials(account, username, keyfile)
				if err != nil {
					panic_the_err(err)
				}

				if !c.Bool("force") {
					err = creds.ValidateCredentials(account, username)
					if err != nil {
						panic_the_err(err)
					}
				}
				cred.Display(os.Stdout)
			},
		},

		{
			Name:  "current",
			Usage: "Show the username and alias of the currently-loaded credentials",
			Action: func(c *cli.Context) {
				AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
				AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
				if AWSAccessKeyId == "" || AWSSecretAccessKey == "" {
					err := errors.New("No amazon credentials are currently in your environment")
					panic_the_err(err)
				}
				cred := Credential{
					KeyId:     AWSAccessKeyId,
					SecretKey: AWSSecretAccessKey,
				}
				username, alias, err := getAWSUsernameAndAlias(cred)
				if err != nil {
					panic_the_err(err)
				}
				fmt.Printf("%s@%s\n", username, alias)
			},
		},

		{
			Name:  "display",
			Usage: "Display loaded AWS credentials",
			Action: func(c *cli.Context) {
				AWSAccessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
				AWSSecretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
				fmt.Printf("AWS_ACCESS_KEY_ID: %s\n", AWSAccessKeyId)
				fmt.Printf("AWS_SECRET_ACCESS_KEY: %s\n", AWSSecretAccessKey)
			},
		},

		{
			Name:  "list",
			Usage: "List available AWS credentials",
			Action: func(c *cli.Context) {
				rootDir, err := os.Open(getRootPath())
				if err != nil {
					panic_the_err(err)
				}
				set, err := listAvailableCredentials(rootDir)
				if err != nil {
					panic_the_err(err)
				}
				for _, cred := range set {
					fmt.Println(cred)
				}
			},
		},

		{
			Name:  "rotate",
			Usage: "Rotate current AWS credentials, deleting the oldest",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "lifetime, l",
					Value: 0,
					Usage: "\n        New credential lifetime in seconds (0 means forever)",
				},
				cli.StringSliceFlag{
					Name:  "key, k",
					Value: &cli.StringSlice{},
					Usage: "\n        SSH public keys for encryption",
				},
				cli.StringSliceFlag{
					Name:  "env, e",
					Value: &cli.StringSlice{},
					Usage: "\n        Environment variables to set in the form VAR=value",
				},
			},
			Action: func(c *cli.Context) {
				cred, _, _, pubkeys, lifetime, err := parseSaveArgs(c)
				panic_the_err(err)
				username, account, err := getAWSUsernameAndAlias(cred)
				panic_the_err(err)
				err = (&cred).rotateCredentials(username)
				panic_the_err(err)
				err = SaveCredentials(cred, username, account, pubkeys, lifetime, c.Bool("force"))
				panic_the_err(err)
			},
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package patch

import (
	"fmt"
	"io/ioutil"
	"reflect"
	"strings"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var patchTypes = map[string]types.PatchType{"json": types.JSONPatchType, "merge": types.MergePatchType, "strategic": types.StrategicMergePatchType}

// PatchOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type PatchOptions struct {
	resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	PrintFlags  *genericclioptions.PrintFlags
	ToPrinter   func(string) (printers.ResourcePrinter, error)
	Recorder    genericclioptions.Recorder

	Local     bool
	PatchType string
	Patch     string
	PatchFile string

	namespace                    string
	enforceNamespace             bool
	dryRunStrategy               cmdutil.DryRunStrategy
	dryRunVerifier               *resource.DryRunVerifier
	outputFormat                 string
	args                         []string
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)
	fieldManager                 string

	genericclioptions.IOStreams
}

var (
	patchLong = templates.LongDesc(i18n.T(`
		Update field(s) of a resource using strategic merge patch, a JSON merge patch, or a JSON patch.

		JSON and YAML formats are accepted.`))

	patchExample = templates.Examples(i18n.T(`
		# Partially update a node using a strategic merge patch. Specify the patch as JSON.
		kubectl patch node k8s-node-1 -p '{"spec":{"unschedulable":true}}'

		# Partially update a node using a strategic merge patch. Specify the patch as YAML.
		kubectl patch node k8s-node-1 -p $'spec:\n unschedulable: true'

		# Partially update a node identified by the type and name specified in "node.json" using strategic merge patch.
		kubectl patch -f node.json -p '{"spec":{"unschedulable":true}}'

		# Update a container's image; spec.containers[*].name is required because it's a merge key.
		kubectl patch pod valid-pod -p '{"spec":{"containers":[{"name":"kubernetes-serve-hostname","image":"new image"}]}}'

		# Update a container's image using a json patch with positional arrays.
		kubectl patch pod valid-pod --type='json' -p='[{"op": "replace", "path": "/spec/containers/0/image", "value":"new image"}]'`))
)

func NewPatchOptions(ioStreams genericclioptions.IOStreams) *PatchOptions {
	return &PatchOptions{
		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		PrintFlags:  genericclioptions.NewPrintFlags("patched").WithTypeSetter(scheme.Scheme),
		IOStreams:   ioStreams,
	}
}

func NewCmdPatch(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewPatchOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "patch (-f FILENAME | TYPE NAME) [-p PATCH|--patch-file FILE]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update field(s) of a resource"),
		Long:                  patchLong,
		Example:               patchExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunPatch())
		},
	}

	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().StringVarP(&o.Patch, "patch", "p", "", "The patch to be applied to the resource JSON file.")
	cmd.Flags().StringVar(&o.PatchFile, "patch-file", "", "A file containing a patch to be applied to the resource.")
	cmd.Flags().StringVar(&o.PatchType, "type", "strategic", fmt.Sprintf("The type of patch being provided; one of %v", sets.StringKeySet(patchTypes).List()))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, "identifying the resource to update")
	cmd.Flags().BoolVar(&o.Local, "local", o.Local, "If true, patch will operate on the content of the file, not the server-side resource.")
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-patch")

	return cmd
}

func (o *PatchOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}

	cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)
	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation

		return o.PrintFlags.ToPrinter()
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.args = args
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, discoveryClient)

	return nil
}

func (o *PatchOptions) Validate() error {
	if len(o.Patch) > 0 && len(o.PatchFile) > 0 {
		return fmt.Errorf("cannot specify --patch and --patch-file together")
	}
	if len(o.Patch) == 0 && len(o.PatchFile) == 0 {
		return fmt.Errorf("must specify --patch or --patch-file containing the contents of the patch")
	}
	if o.Local && len(o.args) != 0 {
		return fmt.Errorf("cannot specify --local and server resources")
	}
	if o.Local && o.dryRunStrategy == cmdutil.DryRunServer {
		return fmt.Errorf("cannot specify --local and --dry-run=server - did you mean --dry-run=client?")
	}
	if len(o.PatchType) != 0 {
		if _, ok := patchTypes[strings.ToLower(o.PatchType)]; !ok {
			return fmt.Errorf("--type must be one of %v, not %q", sets.StringKeySet(patchTypes).List(), o.PatchType)
		}
	}

	return nil
}

func (o *PatchOptions) RunPatch() error {
	patchType := types.StrategicMergePatchType
	if len(o.PatchType) != 0 {
		patchType = patchTypes[strings.ToLower(o.PatchType)]
	}

	var patchBytes []byte
	if len(o.PatchFile) > 0 {
		var err error
		patchBytes, err = ioutil.ReadFile(o.PatchFile)
		if err != nil {
			return fmt.Errorf("unable to read patch file: %v", err)
		}
	} else {
		patchBytes = []byte(o.Patch)
	}

	patchBytes, err := yaml.ToJSON(patchBytes)
	if err != nil {
		return fmt.Errorf("unable to parse %q: %v", o.Patch, err)
	}

	r := o.builder.
		Unstructured().
		ContinueOnError().
		LocalParam(o.Local).
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	err = r.Err()
	if err != nil {
		return err
	}

	count := 0
	err = r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		count++
		name, namespace := info.Name, info.Namespace

		if !o.Local && o.dryRunStrategy != cmdutil.DryRunClient {
			mapping := info.ResourceMapping()
			if o.dryRunStrategy == cmdutil.DryRunServer {
				if err := o.dryRunVerifier.HasSupport(mapping.GroupVersionKind); err != nil {
					return err
				}
			}
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}

			helper := resource.
				NewHelper(client, mapping).
				DryRun(o.dryRunStrategy == cmdutil.DryRunServer).
				WithFieldManager(o.fieldManager)
			patchedObj, err := helper.Patch(namespace, name, patchType, patchBytes, nil)
			if err != nil {
				return err
			}

			didPatch := !reflect.DeepEqual(info.Object, patchedObj)

			// if the recorder makes a change, compute and create another patch
			if mergePatch, err := o.Recorder.MakeRecordMergePatch(patchedObj); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			} else if len(mergePatch) > 0 {
				if recordedObj, err := helper.Patch(info.Namespace, info.Name, types.MergePatchType, mergePatch, nil); err != nil {
					klog.V(4).Infof("error recording reason: %v", err)
				} else {
					patchedObj = recordedObj
				}
			}

			printer, err := o.ToPrinter(patchOperation(didPatch))
			if err != nil {
				return err
			}
			return printer.PrintObj(patchedObj, o.Out)
		}

		originalObjJS, err := runtime.Encode(unstructured.UnstructuredJSONScheme, info.Object)
		if err != nil {
			return err
		}

		originalPatchedObjJS, err := getPatchedJSON(patchType, originalObjJS, patchBytes, info.Object.GetObjectKind().GroupVersionKind(), scheme.Scheme)
		if err != nil {
			return err
		}

		targetObj, err := runtime.Decode(unstructured.UnstructuredJSONScheme, originalPatchedObjJS)
		if err != nil {
			return err
		}

		didPatch := !reflect.DeepEqual(info.Object, targetObj)
		printer, err := o.ToPrinter(patchOperation(didPatch))
		if err != nil {
			return err
		}
		return printer.PrintObj(targetObj, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to patch")
	}
	return nil
}

func getPatchedJSON(patchType types.PatchType, originalJS, patchJS []byte, gvk schema.GroupVersionKind, creater runtime.ObjectCreater) ([]byte, error) {
	switch patchType {
	case types.JSONPatchType:
		patchObj, err := jsonpatch.DecodePatch(patchJS)
		if err != nil {
			return nil, err
		}
		bytes, err := patchObj.Apply(originalJS)
		// TODO: This is pretty hacky, we need a better structured error from the json-patch
		if err != nil && strings.Contains(err.Error(), "doc is missing key") {
			msg := err.Error()
			ix := strings.Index(msg, "key:")
			key := msg[ix+5:]
			return bytes, fmt.Errorf("Object to be patched is missing field (%s)", key)
		}
		return bytes, err

	case types.MergePatchType:
		return jsonpatch.MergePatch(originalJS, patchJS)

	case types.StrategicMergePatchType:
		// get a typed object for this GVK if we need to apply a strategic merge patch
		obj, err := creater.New(gvk)
		if err != nil {
			return nil, fmt.Errorf("cannot apply strategic merge patch for %s locally, try --type merge", gvk.String())
		}
		return strategicpatch.StrategicMergePatch(originalJS, patchJS, obj)

	default:
		// only here as a safety net - go-restful filters content-type
		return nil, fmt.Errorf("unknown Content-Type header for patch: %v", patchType)
	}
}

func patchOperation(didPatch bool) string {
	if didPatch {
		return "patched"
	}
	return "patched (no change)"
}
<commit_msg>kubectl: fix error message for client-side SMP for CRs<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package patch

import (
	"fmt"
	"io/ioutil"
	"reflect"
	"strings"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var patchTypes = map[string]types.PatchType{"json": types.JSONPatchType, "merge": types.MergePatchType, "strategic": types.StrategicMergePatchType}

// PatchOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type PatchOptions struct {
	resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	PrintFlags  *genericclioptions.PrintFlags
	ToPrinter   func(string) (printers.ResourcePrinter, error)
	Recorder    genericclioptions.Recorder

	Local     bool
	PatchType string
	Patch     string
	PatchFile string

	namespace                    string
	enforceNamespace             bool
	dryRunStrategy               cmdutil.DryRunStrategy
	dryRunVerifier               *resource.DryRunVerifier
	outputFormat                 string
	args                         []string
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)
	fieldManager                 string

	genericclioptions.IOStreams
}

var (
	patchLong = templates.LongDesc(i18n.T(`
		Update field(s) of a resource using strategic merge patch, a JSON merge patch, or a JSON patch.

		JSON and YAML formats are accepted.`))

	patchExample = templates.Examples(i18n.T(`
		# Partially update a node using a strategic merge patch. Specify the patch as JSON.
		kubectl patch node k8s-node-1 -p '{"spec":{"unschedulable":true}}'

		# Partially update a node using a strategic merge patch. Specify the patch as YAML.
		kubectl patch node k8s-node-1 -p $'spec:\n unschedulable: true'

		# Partially update a node identified by the type and name specified in "node.json" using strategic merge patch.
		kubectl patch -f node.json -p '{"spec":{"unschedulable":true}}'

		# Update a container's image; spec.containers[*].name is required because it's a merge key.
		kubectl patch pod valid-pod -p '{"spec":{"containers":[{"name":"kubernetes-serve-hostname","image":"new image"}]}}'

		# Update a container's image using a json patch with positional arrays.
		kubectl patch pod valid-pod --type='json' -p='[{"op": "replace", "path": "/spec/containers/0/image", "value":"new image"}]'`))
)

func NewPatchOptions(ioStreams genericclioptions.IOStreams) *PatchOptions {
	return &PatchOptions{
		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		PrintFlags:  genericclioptions.NewPrintFlags("patched").WithTypeSetter(scheme.Scheme),
		IOStreams:   ioStreams,
	}
}

func NewCmdPatch(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewPatchOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "patch (-f FILENAME | TYPE NAME) [-p PATCH|--patch-file FILE]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update field(s) of a resource"),
		Long:                  patchLong,
		Example:               patchExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunPatch())
		},
	}

	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().StringVarP(&o.Patch, "patch", "p", "", "The patch to be applied to the resource JSON file.")
	cmd.Flags().StringVar(&o.PatchFile, "patch-file", "", "A file containing a patch to be applied to the resource.")
	cmd.Flags().StringVar(&o.PatchType, "type", "strategic", fmt.Sprintf("The type of patch being provided; one of %v", sets.StringKeySet(patchTypes).List()))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, "identifying the resource to update")
	cmd.Flags().BoolVar(&o.Local, "local", o.Local, "If true, patch will operate on the content of the file, not the server-side resource.")
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-patch")

	return cmd
}

func (o *PatchOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}

	cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)
	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation

		return o.PrintFlags.ToPrinter()
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.args = args
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, discoveryClient)

	return nil
}

func (o *PatchOptions) Validate() error {
	if len(o.Patch) > 0 && len(o.PatchFile) > 0 {
		return fmt.Errorf("cannot specify --patch and --patch-file together")
	}
	if len(o.Patch) == 0 && len(o.PatchFile) == 0 {
		return fmt.Errorf("must specify --patch or --patch-file containing the contents of the patch")
	}
	if o.Local && len(o.args) != 0 {
		return fmt.Errorf("cannot specify --local and server resources")
	}
	if o.Local && o.dryRunStrategy == cmdutil.DryRunServer {
		return fmt.Errorf("cannot specify --local and --dry-run=server - did you mean --dry-run=client?")
	}
	if len(o.PatchType) != 0 {
		if _, ok := patchTypes[strings.ToLower(o.PatchType)]; !ok {
			return fmt.Errorf("--type must be one of %v, not %q", sets.StringKeySet(patchTypes).List(), o.PatchType)
		}
	}

	return nil
}

func (o *PatchOptions) RunPatch() error {
	patchType := types.StrategicMergePatchType
	if len(o.PatchType) != 0 {
		patchType = patchTypes[strings.ToLower(o.PatchType)]
	}

	var patchBytes []byte
	if len(o.PatchFile) > 0 {
		var err error
		patchBytes, err = ioutil.ReadFile(o.PatchFile)
		if err != nil {
			return fmt.Errorf("unable to read patch file: %v", err)
		}
	} else {
		patchBytes = []byte(o.Patch)
	}

	patchBytes, err := yaml.ToJSON(patchBytes)
	if err != nil {
		return fmt.Errorf("unable to parse %q: %v", o.Patch, err)
	}

	r := o.builder.
		Unstructured().
		ContinueOnError().
		LocalParam(o.Local).
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	err = r.Err()
	if err != nil {
		return err
	}

	count := 0
	err = r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		count++
		name, namespace := info.Name, info.Namespace

		if !o.Local && o.dryRunStrategy != cmdutil.DryRunClient {
			mapping := info.ResourceMapping()
			if o.dryRunStrategy == cmdutil.DryRunServer {
				if err := o.dryRunVerifier.HasSupport(mapping.GroupVersionKind); err != nil {
					return err
				}
			}
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}

			helper := resource.
				NewHelper(client, mapping).
				DryRun(o.dryRunStrategy == cmdutil.DryRunServer).
				WithFieldManager(o.fieldManager)
			patchedObj, err := helper.Patch(namespace, name, patchType, patchBytes, nil)
			if err != nil {
				return err
			}

			didPatch := !reflect.DeepEqual(info.Object, patchedObj)

			// if the recorder makes a change, compute and create another patch
			if mergePatch, err := o.Recorder.MakeRecordMergePatch(patchedObj); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			} else if len(mergePatch) > 0 {
				if recordedObj, err := helper.Patch(info.Namespace, info.Name, types.MergePatchType, mergePatch, nil); err != nil {
					klog.V(4).Infof("error recording reason: %v", err)
				} else {
					patchedObj = recordedObj
				}
			}

			printer, err := o.ToPrinter(patchOperation(didPatch))
			if err != nil {
				return err
			}
			return printer.PrintObj(patchedObj, o.Out)
		}

		originalObjJS, err := runtime.Encode(unstructured.UnstructuredJSONScheme, info.Object)
		if err != nil {
			return err
		}

		originalPatchedObjJS, err := getPatchedJSON(patchType, originalObjJS, patchBytes, info.Object.GetObjectKind().GroupVersionKind(), scheme.Scheme)
		if err != nil {
			return err
		}

		targetObj, err := runtime.Decode(unstructured.UnstructuredJSONScheme, originalPatchedObjJS)
		if err != nil {
			return err
		}

		didPatch := !reflect.DeepEqual(info.Object, targetObj)
		printer, err := o.ToPrinter(patchOperation(didPatch))
		if err != nil {
			return err
		}
		return printer.PrintObj(targetObj, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to patch")
	}
	return nil
}

func getPatchedJSON(patchType types.PatchType, originalJS, patchJS []byte, gvk schema.GroupVersionKind, creater runtime.ObjectCreater) ([]byte, error) {
	switch patchType {
	case types.JSONPatchType:
		patchObj, err := jsonpatch.DecodePatch(patchJS)
		if err != nil {
			return nil, err
		}
		bytes, err := patchObj.Apply(originalJS)
		// TODO: This is pretty hacky, we need a better structured error from the json-patch
		if err != nil && strings.Contains(err.Error(), "doc is missing key") {
			msg := err.Error()
			ix := strings.Index(msg, "key:")
			key := msg[ix+5:]
			return bytes, fmt.Errorf("Object to be patched is missing field (%s)", key)
		}
		return bytes, err

	case types.MergePatchType:
		return jsonpatch.MergePatch(originalJS, patchJS)

	case types.StrategicMergePatchType:
		// get a typed object for this GVK if we need to apply a strategic merge patch
		obj, err := creater.New(gvk)
		if err != nil {
			return nil, fmt.Errorf("strategic merge patch is not supported for %s locally, try --type merge", gvk.String())
		}
		return strategicpatch.StrategicMergePatch(originalJS, patchJS, obj)

	default:
		// only here as a safety net - go-restful filters content-type
		return nil, fmt.Errorf("unknown Content-Type header for patch: %v", patchType)
	}
}

func patchOperation(didPatch bool) string {
	if didPatch {
		return "patched"
	}
	return "patched (no change)"
}
<|endoftext|>


<commit_before>package connpool

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/sys/unix"

	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	rpc "github.com/telepresenceio/telepresence/rpc/v2/manager"
)

// The idleDuration controls how long a dialer for a specific proto+from-to address combination remains alive without
// reading or writing any messages. The dialer is normally closed by one of the peers.
const connTTL = time.Minute
const dialTimeout = 30 * time.Second

const (
	notConnected = int32(iota)
	halfConnected
	connected
)

type TunnelStream interface {
	Send(*rpc.ConnMessage) error
	Recv() (*rpc.ConnMessage, error)
}

// The dialer takes care of dispatching messages between gRPC and UDP connections
type dialer struct {
	id            ConnID
	release       func()
	bidiStream    TunnelStream
	incoming      chan Message
	conn          net.Conn
	idleTimer     *time.Timer
	idleLock      sync.Mutex
	connected     int32
	writerClosing chan struct{}
}

// NewDialer creates a new handler that dispatches messages in both directions between the given gRPC server
// and the destination identified by the given connID.
//
// The handler remains active until it's been idle for idleDuration, at which time it will automatically close
// and call the release function it got from the connpool.Pool to ensure that it gets properly released.
func NewDialer(connID ConnID, bidiStream TunnelStream, release func()) Handler {
	return &dialer{
		id:            connID,
		bidiStream:    bidiStream,
		release:       release,
		incoming:      make(chan Message, 10),
		writerClosing: make(chan struct{}),
		connected:     notConnected,
	}
}

// HandlerFromConn is like NewHandler but initializes the handler with an already existing connection.
func HandlerFromConn(connID ConnID, bidiStream TunnelStream, release func(), conn net.Conn) Handler {
	return &dialer{
		id:            connID,
		bidiStream:    bidiStream,
		release:       release,
		incoming:      make(chan Message, 10),
		writerClosing: make(chan struct{}),
		connected:     halfConnected,
		conn:          conn,
	}
}

func (h *dialer) Start(ctx context.Context) {
	// Set up the idle timer to close and release this handler when it's been idle for a while.
	h.idleTimer = time.NewTimer(connTTL)

	switch h.connected {
	case notConnected:
		if h.id.Protocol() == unix.IPPROTO_UDP {
			h.open(ctx)
		}
	case halfConnected:
		// Connection is created by listener on this side. Establish other
		// side using a control message and start the loops
		h.connected = connected
		h.sendTCD(ctx, Connect)
	}
}

func (h *dialer) open(ctx context.Context) ControlCode {
	if !atomic.CompareAndSwapInt32(&h.connected, notConnected, connected) {
		// already connected
		return ConnectOK
	}
	dlog.Debugf(ctx, "   CONN %s, dialing", h.id)
	dialer := net.Dialer{Timeout: dialTimeout}
	conn, err := dialer.DialContext(ctx, h.id.ProtocolString(), h.id.DestinationAddr().String())
	if err != nil {
		dlog.Errorf(ctx, "%s: failed to establish connection: %v", h.id, err)
		return ConnectReject
	}
	h.conn = conn
	dlog.Debugf(ctx, "   CONN %s, dial answered", h.id)
	go h.writeLoop(ctx)
	go h.readLoop(ctx)
	return ConnectOK
}

func (h *dialer) handleControl(ctx context.Context, cm Control) {
	dlog.Debugf(ctx, "<- GRPC %s", cm)
	switch cm.Code() {
	case Connect:
		h.sendTCD(ctx, h.open(ctx))
	case ConnectOK:
		go h.writeLoop(ctx)
		go h.readLoop(ctx)
	case Disconnect, ConnectReject:
		h.Close(ctx)
		h.sendTCD(ctx, DisconnectOK)
	case ReadClosed:
		// If there's nothing more for us to read, then it's time to ask the writer to close.
		// We do this instead of h.Close() because it will ensure that anything the writeLoop
		// is already dealing with will be written. Will prevent connections from being
		// abruptly terminated after the source has sent everything but before we've had
		// a chance to forward it.
		if atomic.LoadInt32(&h.connected) > 0 {
			select {
			case <-h.writerClosing:
				// Already closed, nothing to do
			default:
				close(h.writerClosing)
			}
		}
	case WriteClosed:
		h.Close(ctx)
	case KeepAlive:
		h.resetIdle()
	default:
		dlog.Errorf(ctx, "%s: unhandled connection control message: %s", h.id, cm)
	}
}

// HandleMessage sends a package to the underlying TCP/UDP connection
func (h *dialer) HandleMessage(ctx context.Context, dg Message) {
	if ctrl, ok := dg.(Control); ok {
		h.handleControl(ctx, ctrl)
		return
	}
	select {
	case <-ctx.Done():
		return
	case h.incoming <- dg:
	}
}

// Close will close the underlying TCP/UDP connection
func (h *dialer) Close(_ context.Context) {
	if atomic.CompareAndSwapInt32(&h.connected, connected, notConnected) {
		h.release()
		if h.conn != nil {
			_ = h.conn.Close()
		}
	}
}

func (h *dialer) sendTCD(ctx context.Context, code ControlCode) {
	ctrl := NewControl(h.id, code, nil)
	dlog.Debugf(ctx, "-> GRPC %s", ctrl)
	err := h.bidiStream.Send(ctrl.TunnelMessage())
	if err != nil {
		dlog.Errorf(ctx, "failed to send control message: %v", err)
	}
}

func (h *dialer) readLoop(ctx context.Context) {
	defer func() {
		// allow write to drain and perform the close of the connection
		dtime.SleepWithContext(ctx, 100*time.Millisecond)
		close(h.incoming)
	}()
	b := make([]byte, 0x8000)
	for ctx.Err() == nil {
		n, err := h.conn.Read(b)
		if err != nil {
			if atomic.LoadInt32(&h.connected) > 0 && ctx.Err() == nil {
				if err != io.EOF {
					dlog.Errorf(ctx, "!! CONN %s, conn read: %v", h.id, err)
				}
			}
			return
		}
		if !h.resetIdle() {
			return
		}
		if n > 0 {
			dlog.Debugf(ctx, "<- CONN %s, len %d", h.id, n)
			if err = h.bidiStream.Send(NewMessage(h.id, b[:n]).TunnelMessage()); err != nil {
				if ctx.Err() == nil {
					dlog.Errorf(ctx, "!! GRPC %s, send: %v", h.id, err)
				}
				return
			}
			dlog.Debugf(ctx, "-> GRPC %s, len %d", h.id, n)
		}
	}
}

func (h *dialer) writeLoop(ctx context.Context) {
	defer h.Close(ctx)
	for {
		select {
		case <-ctx.Done():
			return
		case <-h.idleTimer.C:
			return
		case dg := <-h.incoming:
			if dg == nil {
				// h.incoming was closed by the reader and is now drained.
				if h.id.Protocol() == unix.IPPROTO_TCP {
					h.sendTCD(ctx, ReadClosed)
				}
				return
			}
			if !h.resetIdle() {
				return
			}
			payload := dg.Payload()
			pn := len(payload)
			dlog.Debugf(ctx, "<- GRPC %s, len %d", h.id, pn)
			for n := 0; n < pn; {
				wn, err := h.conn.Write(payload[n:])
				if err != nil {
					if atomic.LoadInt32(&h.connected) > 0 && ctx.Err() == nil {
						if h.id.Protocol() == unix.IPPROTO_TCP {
							h.sendTCD(ctx, WriteClosed)
						}
						dlog.Errorf(ctx, "!! CONN %s, write: %v", h.id, err)
					}
					return
				}
				dlog.Debugf(ctx, "-> CONN %s, len %d", h.id, wn)
				n += wn
			}
		case <-h.writerClosing:
			return
		}
	}
}

func (h *dialer) resetIdle() bool {
	h.idleLock.Lock()
	stopped := h.idleTimer.Stop()
	if stopped {
		h.idleTimer.Reset(connTTL)
	}
	h.idleLock.Unlock()
	return stopped
}
<commit_msg>PR comments<commit_after>package connpool

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/sys/unix"

	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	rpc "github.com/telepresenceio/telepresence/rpc/v2/manager"
)

// The idleDuration controls how long a dialer for a specific proto+from-to address combination remains alive without
// reading or writing any messages. The dialer is normally closed by one of the peers.
const connTTL = time.Minute
const dialTimeout = 30 * time.Second

const (
	notConnected = int32(iota)
	halfConnected
	connected
)

type TunnelStream interface {
	Send(*rpc.ConnMessage) error
	Recv() (*rpc.ConnMessage, error)
}

// The dialer takes care of dispatching messages between gRPC and UDP connections
type dialer struct {
	id            ConnID
	release       func()
	bidiStream    TunnelStream
	incoming      chan Message
	conn          net.Conn
	idleTimer     *time.Timer
	idleLock      sync.Mutex
	connected     int32
	writerClosing chan struct{}
}

// NewDialer creates a new handler that dispatches messages in both directions between the given gRPC server
// and the destination identified by the given connID.
//
// The handler remains active until it's been idle for idleDuration, at which time it will automatically close
// and call the release function it got from the connpool.Pool to ensure that it gets properly released.
func NewDialer(connID ConnID, bidiStream TunnelStream, release func()) Handler {
	return &dialer{
		id:            connID,
		bidiStream:    bidiStream,
		release:       release,
		incoming:      make(chan Message, 10),
		writerClosing: make(chan struct{}),
		connected:     notConnected,
	}
}

// HandlerFromConn is like NewHandler but initializes the handler with an already existing connection.
func HandlerFromConn(connID ConnID, bidiStream TunnelStream, release func(), conn net.Conn) Handler {
	return &dialer{
		id:            connID,
		bidiStream:    bidiStream,
		release:       release,
		incoming:      make(chan Message, 10),
		writerClosing: make(chan struct{}),
		connected:     halfConnected,
		conn:          conn,
	}
}

func (h *dialer) Start(ctx context.Context) {
	// Set up the idle timer to close and release this handler when it's been idle for a while.
	h.idleTimer = time.NewTimer(connTTL)

	switch h.connected {
	case notConnected:
		if h.id.Protocol() == unix.IPPROTO_UDP {
			h.open(ctx)
		}
	case halfConnected:
		// Connection is created by listener on this side. Establish other
		// side using a control message and start the loops
		h.connected = connected
		h.sendTCD(ctx, Connect)
	}
}

func (h *dialer) open(ctx context.Context) ControlCode {
	if !atomic.CompareAndSwapInt32(&h.connected, notConnected, connected) {
		// already connected
		return ConnectOK
	}
	dlog.Debugf(ctx, "   CONN %s, dialing", h.id)
	dialer := net.Dialer{Timeout: dialTimeout}
	conn, err := dialer.DialContext(ctx, h.id.ProtocolString(), h.id.DestinationAddr().String())
	if err != nil {
		dlog.Errorf(ctx, "%s: failed to establish connection: %v", h.id, err)
		return ConnectReject
	}
	h.conn = conn
	dlog.Debugf(ctx, "   CONN %s, dial answered", h.id)
	go h.writeLoop(ctx)
	go h.readLoop(ctx)
	return ConnectOK
}

func (h *dialer) handleControl(ctx context.Context, cm Control) {
	dlog.Debugf(ctx, "<- GRPC %s", cm)
	switch cm.Code() {
	case Connect:
		h.sendTCD(ctx, h.open(ctx))
	case ConnectOK:
		go h.writeLoop(ctx)
		go h.readLoop(ctx)
	case Disconnect:
		h.Close(ctx)
		h.sendTCD(ctx, DisconnectOK)
	case ReadClosed:
		// If there's nothing more for us to read, then it's time to ask the writer to close.
		// We do this instead of h.Close() because it will ensure that anything the writeLoop
		// is already dealing with will be written. Will prevent connections from being
		// abruptly terminated after the source has sent everything but before we've had
		// a chance to forward it.
		if atomic.LoadInt32(&h.connected) > 0 {
			select {
			case <-h.writerClosing:
				// Already closed, nothing to do
			default:
				close(h.writerClosing)
			}
		}
	case WriteClosed, ConnectReject:
		h.Close(ctx)
	case KeepAlive:
		h.resetIdle()
	default:
		dlog.Errorf(ctx, "%s: unhandled connection control message: %s", h.id, cm)
	}
}

// HandleMessage sends a package to the underlying TCP/UDP connection
func (h *dialer) HandleMessage(ctx context.Context, dg Message) {
	if ctrl, ok := dg.(Control); ok {
		h.handleControl(ctx, ctrl)
		return
	}
	select {
	case <-ctx.Done():
		return
	case h.incoming <- dg:
	}
}

// Close will close the underlying TCP/UDP connection
func (h *dialer) Close(_ context.Context) {
	if atomic.CompareAndSwapInt32(&h.connected, connected, notConnected) {
		h.release()
		if h.conn != nil {
			_ = h.conn.Close()
		}
	}
}

func (h *dialer) sendTCD(ctx context.Context, code ControlCode) {
	ctrl := NewControl(h.id, code, nil)
	dlog.Debugf(ctx, "-> GRPC %s", ctrl)
	err := h.bidiStream.Send(ctrl.TunnelMessage())
	if err != nil {
		dlog.Errorf(ctx, "failed to send control message: %v", err)
	}
}

func (h *dialer) readLoop(ctx context.Context) {
	defer func() {
		// allow write to drain and perform the close of the connection
		dtime.SleepWithContext(ctx, 100*time.Millisecond)
		close(h.incoming)
	}()
	b := make([]byte, 0x8000)
	for ctx.Err() == nil {
		n, err := h.conn.Read(b)
		if err != nil {
			if atomic.LoadInt32(&h.connected) > 0 && ctx.Err() == nil {
				if err != io.EOF {
					dlog.Errorf(ctx, "!! CONN %s, conn read: %v", h.id, err)
				}
			}
			return
		}
		if !h.resetIdle() {
			return
		}
		if n > 0 {
			dlog.Debugf(ctx, "<- CONN %s, len %d", h.id, n)
			if err = h.bidiStream.Send(NewMessage(h.id, b[:n]).TunnelMessage()); err != nil {
				if ctx.Err() == nil {
					dlog.Errorf(ctx, "!! GRPC %s, send: %v", h.id, err)
				}
				return
			}
			dlog.Debugf(ctx, "-> GRPC %s, len %d", h.id, n)
		}
	}
}

func (h *dialer) writeLoop(ctx context.Context) {
	defer h.Close(ctx)
	for {
		select {
		case <-ctx.Done():
			return
		case <-h.idleTimer.C:
			return
		case dg := <-h.incoming:
			if dg == nil {
				// h.incoming was closed by the reader and is now drained.
				if h.id.Protocol() == unix.IPPROTO_TCP {
					h.sendTCD(ctx, ReadClosed)
				}
				return
			}
			if !h.resetIdle() {
				return
			}
			payload := dg.Payload()
			pn := len(payload)
			dlog.Debugf(ctx, "<- GRPC %s, len %d", h.id, pn)
			for n := 0; n < pn; {
				wn, err := h.conn.Write(payload[n:])
				if err != nil {
					if atomic.LoadInt32(&h.connected) > 0 && ctx.Err() == nil {
						if h.id.Protocol() == unix.IPPROTO_TCP {
							h.sendTCD(ctx, WriteClosed)
						}
						dlog.Errorf(ctx, "!! CONN %s, write: %v", h.id, err)
					}
					return
				}
				dlog.Debugf(ctx, "-> CONN %s, len %d", h.id, wn)
				n += wn
			}
		case <-h.writerClosing:
			return
		}
	}
}

func (h *dialer) resetIdle() bool {
	h.idleLock.Lock()
	stopped := h.idleTimer.Stop()
	if stopped {
		h.idleTimer.Reset(connTTL)
	}
	h.idleLock.Unlock()
	return stopped
}
<|endoftext|>


<commit_before>package editcmd

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/skatteetaten/aoc/pkg/auroraconfig"
	"github.com/skatteetaten/aoc/pkg/cmdoptions"
	"github.com/skatteetaten/aoc/pkg/configuration"
	"github.com/skatteetaten/aoc/pkg/fileutil"
	"github.com/skatteetaten/aoc/pkg/jsonutil"
	"github.com/skatteetaten/aoc/pkg/serverapi"
	"io/ioutil"
	"os"
	"strings"
)

const usageString = "Usage: aoc edit file [env/]<filename> | secret <vaultname> <secretname> "
const secretUseageString = "Usage: aoc edit secret <vaultname> <secretname>"
const fileUseageString = "Usage: aoc edit file [env/]<filename>"

const commentString = "# "
const editMessage = `# Please edit the object below. Lines beginning with a '#' will be ignored,
# and an empty file will abort the edit. If an error occurs while saving this file will be
# reopened with the relevant failures.
#
`

type EditcmdClass struct {
	configuration configuration.ConfigurationClass
}

func (editcmdClass *EditcmdClass) getAffiliation() (affiliation string) {
	if editcmdClass.configuration.GetOpenshiftConfig() != nil {
		affiliation = editcmdClass.configuration.GetOpenshiftConfig().Affiliation
	}
	return
}

func (editcmdClass *EditcmdClass) EditFile(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {

	var filename string = args[1]
	var content string
	var version string

	content, version, err = auroraconfig.GetContent(filename, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	if err != nil {
		return "", err
	}
	content = jsonutil.PrettyPrintJson(content)

	var editCycleDone bool
	var modifiedContent = content

	for editCycleDone == false {
		contentBeforeEdit := modifiedContent
		modifiedContent, err = editString(editMessage + modifiedContent)
		if err != nil {
			return "", err
		}
		if (stripComments(modifiedContent) == stripComments(contentBeforeEdit)) || stripComments(modifiedContent) == stripComments(content) {
			if stripComments(modifiedContent) != stripComments(content) {
				tempfile, err := createTempFile(stripComments(modifiedContent))
				if err != nil {
					return "", nil
				}
				output += "A copy of your changes har been stored to \"" + tempfile + "\"\n"
			}
			output += "Edit cancelled, no valid changes were saved."
			if persistentOptions.Debug {
				fmt.Println("DEBUG: Content of modified file:")
				fmt.Println(modifiedContent)
			}
			return output, nil
		}
		modifiedContent = stripComments(modifiedContent)

		if jsonutil.IsLegalJson(modifiedContent) {
			validationMessages, err := auroraconfig.PutFile(filename, modifiedContent, version, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
			if err != nil {
				if err.Error() == auroraconfig.InvalidConfigurationError {
					modifiedContent, _ = addComments(modifiedContent, validationMessages)
				} else {
					editCycleDone = true
				}
			} else {
				editCycleDone = true
			}
		} else {
			modifiedContent, _ = addComments(modifiedContent, "Illegal JSON Format")
		}
	}

	return
}

func (editcmdClass *EditcmdClass) EditSecret(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {

	var vaultname string = args[1]
	var secretname string = args[2]
	var version string = ""

	secret, version, err := auroraconfig.GetSecret(vaultname, secretname, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	if err != nil {
		return "", err
	}

	var modifiedSecret = secret
	modifiedSecret, err = editString(modifiedSecret)
	if err != nil {
		return "", err
	}

	if modifiedSecret != secret {
		_, err = auroraconfig.PutSecret(vaultname, secretname, modifiedSecret, version, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	}

	return
}

func addComments(content string, comments string) (commentedContent string, err error) {
	var commentLines []string

	const newline = "\n"
	var commentedComments string
	commentLines, _ = contentToLines(comments)
	for lineno := range commentLines {
		commentedComments += commentString + commentLines[lineno] + newline
	}
	commentedContent = commentedComments + content

	return
}

func stripComments(content string) (uncommentedContent string) {
	var contenttLines []string

	var newline = ""
	contenttLines, _ = contentToLines(content)
	for lineno := range contenttLines {
		if strings.TrimLeft(contenttLines[lineno], commentString) == contenttLines[lineno] {
			uncommentedContent += newline + contenttLines[lineno]
			newline = "\n"
		}
	}

	return
}

func contentToLines(content string) (contentLines []string, err error) {

	scanner := bufio.NewScanner(strings.NewReader(content))
	for scanner.Scan() {
		contentLines = append(contentLines, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return
}

func editString(content string) (modifiedContent string, err error) {

	filename, err := createTempFile(content)

	err = fileutil.EditFile(filename)
	if err != nil {
		return
	}

	fileText, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	modifiedContent = string(fileText)

	err = os.Remove(filename)
	if err != nil {
		err = errors.New("WARNING: Unable to delete tempfile " + filename)
	}
	return
}

func createTempFile(content string) (filename string, err error) {
	const tmpFilePrefix = ".aoc_edit_file_"
	var tmpDir = os.TempDir()
	tmpFile, err := ioutil.TempFile(tmpDir, tmpFilePrefix)
	if err != nil {
		return "", errors.New("Unable to create temporary file: " + err.Error())
	}
	if fileutil.IsLegalFileFolder(tmpFile.Name()) != fileutil.SpecIsFile {
		err = errors.New("Internal error: Illegal temp file name: " + tmpFile.Name())
	}
	filename = tmpFile.Name()
	err = ioutil.WriteFile(tmpFile.Name(), []byte(content), 0700)
	if err != nil {
		return
	}
	return
}

func (editcmdClass *EditcmdClass) EditObject(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {
	err = validateEditcmd(args)
	if err != nil {
		return
	}
	if !serverapi.ValidateLogin(editcmdClass.configuration.GetOpenshiftConfig()) {
		return "", errors.New("Not logged in, please use aoc login")
	}

	var commandStr = args[0]
	switch commandStr {
	case "file":
		{
			output, err = editcmdClass.EditFile(args, persistentOptions)
		}
	case "secret":
		{
			output, err = editcmdClass.EditSecret(args, persistentOptions)
		}
	}
	return

}

func validateEditcmd(args []string) (err error) {

	var commandStr = args[0]
	switch commandStr {
	case "file":
		{
			if len(args) != 2 {
				err = errors.New(fileUseageString)
				return
			}
		}
	case "secret":
		{
			if len(args) != 3 {
				err = errors.New(secretUseageString)
				return
			}
		}
	default:
		{
			err = errors.New(usageString)
		}
	}
	return
}
<commit_msg>Added debug of stripcomment<commit_after>package editcmd

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/skatteetaten/aoc/pkg/auroraconfig"
	"github.com/skatteetaten/aoc/pkg/cmdoptions"
	"github.com/skatteetaten/aoc/pkg/configuration"
	"github.com/skatteetaten/aoc/pkg/fileutil"
	"github.com/skatteetaten/aoc/pkg/jsonutil"
	"github.com/skatteetaten/aoc/pkg/serverapi"
	"io/ioutil"
	"os"
	"strings"
)

const usageString = "Usage: aoc edit file [env/]<filename> | secret <vaultname> <secretname> "
const secretUseageString = "Usage: aoc edit secret <vaultname> <secretname>"
const fileUseageString = "Usage: aoc edit file [env/]<filename>"

const commentString = "# "
const editMessage = `# Please edit the object below. Lines beginning with a '#' will be ignored,
# and an empty file will abort the edit. If an error occurs while saving this file will be
# reopened with the relevant failures.
#
`

type EditcmdClass struct {
	configuration configuration.ConfigurationClass
}

func (editcmdClass *EditcmdClass) getAffiliation() (affiliation string) {
	if editcmdClass.configuration.GetOpenshiftConfig() != nil {
		affiliation = editcmdClass.configuration.GetOpenshiftConfig().Affiliation
	}
	return
}

func (editcmdClass *EditcmdClass) EditFile(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {

	var filename string = args[1]
	var content string
	var version string

	content, version, err = auroraconfig.GetContent(filename, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	if err != nil {
		return "", err
	}
	content = jsonutil.PrettyPrintJson(content)

	var editCycleDone bool
	var modifiedContent = content

	for editCycleDone == false {
		contentBeforeEdit := modifiedContent
		modifiedContent, err = editString(editMessage + modifiedContent)
		if err != nil {
			return "", err
		}
		if (stripComments(modifiedContent) == stripComments(contentBeforeEdit)) || stripComments(modifiedContent) == stripComments(content) {
			if stripComments(modifiedContent) != stripComments(content) {
				tempfile, err := createTempFile(stripComments(modifiedContent))
				if err != nil {
					return "", nil
				}
				output += "A copy of your changes har been stored to \"" + tempfile + "\"\n"
			}
			output += "Edit cancelled, no valid changes were saved."
			if persistentOptions.Debug {
				fmt.Println("DEBUG: Content of modified file:")
				fmt.Println(modifiedContent)
				fmt.Println("DEBUG: Content of modified file stripped:")
				fmt.Println(stripComments(modifiedContent))
			}
			return output, nil
		}
		modifiedContent = stripComments(modifiedContent)

		if jsonutil.IsLegalJson(modifiedContent) {
			validationMessages, err := auroraconfig.PutFile(filename, modifiedContent, version, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
			if err != nil {
				if err.Error() == auroraconfig.InvalidConfigurationError {
					modifiedContent, _ = addComments(modifiedContent, validationMessages)
				} else {
					editCycleDone = true
				}
			} else {
				editCycleDone = true
			}
		} else {
			modifiedContent, _ = addComments(modifiedContent, "Illegal JSON Format")
		}
	}

	return
}

func (editcmdClass *EditcmdClass) EditSecret(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {

	var vaultname string = args[1]
	var secretname string = args[2]
	var version string = ""

	secret, version, err := auroraconfig.GetSecret(vaultname, secretname, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	if err != nil {
		return "", err
	}

	var modifiedSecret = secret
	modifiedSecret, err = editString(modifiedSecret)
	if err != nil {
		return "", err
	}

	if modifiedSecret != secret {
		_, err = auroraconfig.PutSecret(vaultname, secretname, modifiedSecret, version, persistentOptions, editcmdClass.getAffiliation(), editcmdClass.configuration.GetOpenshiftConfig())
	}

	return
}

func addComments(content string, comments string) (commentedContent string, err error) {
	var commentLines []string

	const newline = "\n"
	var commentedComments string
	commentLines, _ = contentToLines(comments)
	for lineno := range commentLines {
		commentedComments += commentString + commentLines[lineno] + newline
	}
	commentedContent = commentedComments + content

	return
}

func stripComments(content string) (uncommentedContent string) {
	var contenttLines []string

	var newline = ""
	contenttLines, _ = contentToLines(content)
	for lineno := range contenttLines {
		if strings.TrimLeft(contenttLines[lineno], commentString) == contenttLines[lineno] {
			uncommentedContent += newline + contenttLines[lineno]
			newline = "\n"
		}
	}

	return
}

func contentToLines(content string) (contentLines []string, err error) {

	scanner := bufio.NewScanner(strings.NewReader(content))
	for scanner.Scan() {
		contentLines = append(contentLines, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return
}

func editString(content string) (modifiedContent string, err error) {

	filename, err := createTempFile(content)

	err = fileutil.EditFile(filename)
	if err != nil {
		return
	}

	fileText, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	modifiedContent = string(fileText)

	err = os.Remove(filename)
	if err != nil {
		err = errors.New("WARNING: Unable to delete tempfile " + filename)
	}
	return
}

func createTempFile(content string) (filename string, err error) {
	const tmpFilePrefix = ".aoc_edit_file_"
	var tmpDir = os.TempDir()
	tmpFile, err := ioutil.TempFile(tmpDir, tmpFilePrefix)
	if err != nil {
		return "", errors.New("Unable to create temporary file: " + err.Error())
	}
	if fileutil.IsLegalFileFolder(tmpFile.Name()) != fileutil.SpecIsFile {
		err = errors.New("Internal error: Illegal temp file name: " + tmpFile.Name())
	}
	filename = tmpFile.Name()
	err = ioutil.WriteFile(tmpFile.Name(), []byte(content), 0700)
	if err != nil {
		return
	}
	return
}

func (editcmdClass *EditcmdClass) EditObject(args []string, persistentOptions *cmdoptions.CommonCommandOptions) (output string, err error) {
	err = validateEditcmd(args)
	if err != nil {
		return
	}
	if !serverapi.ValidateLogin(editcmdClass.configuration.GetOpenshiftConfig()) {
		return "", errors.New("Not logged in, please use aoc login")
	}

	var commandStr = args[0]
	switch commandStr {
	case "file":
		{
			output, err = editcmdClass.EditFile(args, persistentOptions)
		}
	case "secret":
		{
			output, err = editcmdClass.EditSecret(args, persistentOptions)
		}
	}
	return

}

func validateEditcmd(args []string) (err error) {

	var commandStr = args[0]
	switch commandStr {
	case "file":
		{
			if len(args) != 2 {
				err = errors.New(fileUseageString)
				return
			}
		}
	case "secret":
		{
			if len(args) != 3 {
				err = errors.New(secretUseageString)
				return
			}
		}
	default:
		{
			err = errors.New(usageString)
		}
	}
	return
}
<|endoftext|>