<commit_before>/*-
 * Copyright 2016 Square Inc.
 *
 * 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 jceks parses JCEKS (Java Cryptogaphy Extension Key Store)
// files and extracts keys and certificates. This module only implements
// a fraction of the JCEKS cryptographic protocols. In particular, it
// implements the SHA1 signature verification of the key store and the
// PBEWithMD5AndDES3CBC cipher for encrypting private keys.
package jceks

import (
	"bytes"
	"crypto"
	"crypto/sha1"
	"crypto/subtle"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/binary"
	"fmt"
	"hash"
	"io"
	"os"
	"time"
)

const (
	jceksMagic   = 0xcececece
	jceksVersion = 0x02
	jksMagic     = 0xfeedfeed
)

var (
	oidKeyProtector = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 42, 2, 17, 1, 1}
	oidPublicKeyRSA = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1}
	oidPublicKeyEC  = asn1.ObjectIdentifier{1, 2, 840, 10045, 2, 1}
)

type encryptedPrivateKeyInfo struct {
	Algo         pkix.AlgorithmIdentifier
	EncryptedKey []byte
}

type privateKeyInfo struct {
	Version    int
	Algo       pkix.AlgorithmIdentifier
	PrivateKey []byte
}

type privateKeyEntry struct {
	date       time.Time
	encodedKey []byte
	certs      []*x509.Certificate
}

// From https://golang.org/src/crypto/x509/sec1.go (also: see RFC 5915)
type ecPrivateKey struct {
	Version       int
	PrivateKey    []byte
	NamedCurveOID asn1.ObjectIdentifier `asn1:"optional,explicit,tag:0"`
	PublicKey     asn1.BitString        `asn1:"optional,explicit,tag:1"`
}

func (e *privateKeyEntry) String() string {
	return fmt.Sprintf("private-key: %s", e.date)
}

func (e *privateKeyEntry) Recover(password []byte) (crypto.PrivateKey, error) {
	var eKey encryptedPrivateKeyInfo
	_, err := asn1.Unmarshal(e.encodedKey, &eKey)
	if err != nil {
		return nil, err
	}
	var decryptedKey []byte
	switch {
	case eKey.Algo.Algorithm.Equal(oidPBEWithMD5AndDES3CBC):
		decryptedKey, err = recoverPBEWithMD5AndDES3CBC(eKey.Algo, eKey.EncryptedKey, password)
	case eKey.Algo.Algorithm.Equal(oidKeyProtector):
		// JavaSoft proprietary key-protection algorithm (used to protect
		// private keys in the keystore implementation that comes with JDK
		// 1.2). We shouldn't need this.
		fallthrough
	default:
		return nil, fmt.Errorf("unsupported encrypted-private-key algorithm: %v", eKey.Algo.Algorithm)
	}
	if err != nil {
		return nil, err
	}

	var pKey privateKeyInfo
	if _, err := asn1.Unmarshal(decryptedKey, &pKey); err != nil {
		return nil, err
	}
	if pKey.Algo.Algorithm.Equal(oidPublicKeyRSA) {
		return x509.ParsePKCS1PrivateKey(pKey.PrivateKey)
	}
	if pKey.Algo.Algorithm.Equal(oidPublicKeyEC) {
		// In a JCEKS file, the EC private key blob contains only the key itself, without the
		// named curve OID. Instead, the named curve OID is in a separate field (in the algorithm
		// indentifier for the keystore entry). However to parse the EC key properly we need the
		// EC key blob to have the curve OID...
		key := ecPrivateKey{}
		oid := asn1.ObjectIdentifier{}
		// Parse EC private key
		_, err := asn1.Unmarshal(pKey.PrivateKey, &key)
		if err != nil {
			return nil, fmt.Errorf("problem parsing ec key asn.1 struct: %s", err)
		}
		// Parse named curve OID from algorithm identifier
		_, err = asn1.Unmarshal(pKey.Algo.Parameters.FullBytes, &oid)
		if err != nil {
			return nil, fmt.Errorf("problem parsing ec key asn.1 struct: %s", err)
		}
		// Update key to add named curve info, re-marshal, and parse
		key.NamedCurveOID = oid
		raw, _ := asn1.Marshal(key)
		return x509.ParseECPrivateKey(raw)
	}
	return nil, fmt.Errorf("unsupported private-key algorithm: %v", pKey.Algo.Algorithm)
}

type trustedCertEntry struct {
	date time.Time
	cert *x509.Certificate
}

func (e *trustedCertEntry) String() string {
	return fmt.Sprintf("trusted-cert: %s", e.date)
}

// KeyStore maintains a map from alias name to the entry for that
// alias. Entries are currently either privateKeyEntry or
// trustedCertEntry.
type KeyStore struct {
	entries map[string]interface{}
}

// readUTF reads a java encoded UTF-8 string. The encoding provides a
// 2-byte prefix indicating the length of the string.
func readUTF(r io.Reader) (string, error) {
	var length uint16
	err := binary.Read(r, binary.BigEndian, &length)
	if err != nil {
		return "", err
	}
	buf := make([]byte, length)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return "", err
	}
	return string(buf), nil
}

// readBytes reads a byte array from the reader. The encoding provides
// a 4-byte prefix indicating the number of bytes which follow.
func readBytes(r io.Reader) ([]byte, error) {
	length, err := readInt32(r)
	if err != nil {
		return nil, err
	}
	buf := make([]byte, length)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	return buf, nil
}

func readInt32(r io.Reader) (int32, error) {
	var v int32
	err := binary.Read(r, binary.BigEndian, &v)
	return v, err
}

func readUint32(r io.Reader) (uint32, error) {
	var v uint32
	err := binary.Read(r, binary.BigEndian, &v)
	return v, err
}

func readDate(r io.Reader) (time.Time, error) {
	var v int64
	err := binary.Read(r, binary.BigEndian, &v)
	if err != nil {
		return time.Time{}, err
	}
	sec := v / 1000
	nsec := (v - sec*1000) * 1000 * 1000
	return time.Unix(sec, nsec), nil
}

// Returns a SHA1 hash which has been pre-keyed with the specified
// password according to the JCEKS algorithm.
func getPreKeyedHash(password []byte) hash.Hash {
	md := sha1.New()
	buf := make([]byte, len(password)*2)
	for i := 0; i < len(password); i++ {
		buf[i*2+1] = password[i]
	}
	md.Write(buf)
	// Yes, "Mighty Aprhodite" is a constant used by this method.
	md.Write([]byte("Mighty Aphrodite"))
	return md
}

func parseHeader(r io.Reader) (uint32, error) {
	magic, err := readUint32(r)
	if err != nil {
		return 0, err
	}
	if magic != jceksMagic && magic != jksMagic {
		return 0, fmt.Errorf("unexpected magic: %08x != (%08x || %08x)", magic, uint32(jceksMagic), uint32(jksMagic))
	}
	version, err := readUint32(r)
	if err != nil {
		return 0, err
	}
	return version, nil
}

func (ks *KeyStore) parsePrivateKey(r io.Reader) error {
	alias, err := readUTF(r)
	if err != nil {
		return err
	}
	entry := &privateKeyEntry{
		certs: []*x509.Certificate{},
	}
	entry.date, err = readDate(r)
	if err != nil {
		return err
	}
	entry.encodedKey, err = readBytes(r)
	if err != nil {
		return err
	}
	nCerts, err := readInt32(r)
	if err != nil {
		return err
	}
	for j := 0; j < int(nCerts); j++ {
		certType, err := readUTF(r)
		if err != nil {
			return err
		}
		if certType != "X.509" {
			return fmt.Errorf("unable to handle certificate type: %s", certType)
		}
		certBytes, err := readBytes(r)
		if err != nil {
			return err
		}
		cert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			return nil
		}
		entry.certs = append(entry.certs, cert)
	}
	ks.entries[alias] = entry
	return nil
}

func (ks *KeyStore) parseTrustedCert(r io.Reader) error {
	alias, err := readUTF(r)
	if err != nil {
		return err
	}
	entry := &trustedCertEntry{}
	entry.date, err = readDate(r)
	if err != nil {
		return err
	}

	certType, err := readUTF(r)
	if err != nil {
		return err
	}
	if certType != "X.509" {
		return fmt.Errorf("unable to handle certificate type: %s", certType)
	}
	certBytes, err := readBytes(r)
	if err != nil {
		return err
	}
	entry.cert, err = x509.ParseCertificate(certBytes)
	if err != nil {
		return nil
	}

	ks.entries[alias] = entry
	return nil
}

// Parse parses the key store from the specified reader.
func (ks *KeyStore) Parse(r io.Reader, password []byte) error {
	var md hash.Hash
	if password != nil {
		md = getPreKeyedHash(password)
		r = io.TeeReader(r, md)
	}

	version, err := parseHeader(r)
	if err != nil {
		return err
	}
	if version != jceksVersion {
		return fmt.Errorf("unexpected version: %d != %d", version, jceksVersion)
	}

	count, err := readInt32(r)
	if err != nil {
		return err
	}
	for i := 0; i < int(count); i++ {
		tag, err := readInt32(r)
		if err != nil {
			return err
		}
		switch tag {
		case 1:
			// Private-key entry
			err := ks.parsePrivateKey(r)
			if err != nil {
				return err
			}
		case 2:
			// Trusted-cert entry
			err := ks.parseTrustedCert(r)
			if err != nil {
				return err
			}
		case 3:
			// Secret-key entry
			return fmt.Errorf("unimplemented: Secret-key")
		default:
			return fmt.Errorf("unimplemented tag: %d", tag)
		}
	}

	if md != nil {
		computed := md.Sum([]byte{})
		actual := make([]byte, len(computed))
		_, err := io.ReadFull(r, actual)
		if err != nil {
			return err
		}
		if subtle.ConstantTimeCompare(computed, actual) != 1 {
			return fmt.Errorf("keystore was tampered with or password was incorrect")
		}
	}

	return nil
}

// GetPrivateKeyAndCerts retrieves the specified private key. Returns
// nil if the private key does not exist or alias points to a non
// private key entry.
func (ks *KeyStore) GetPrivateKeyAndCerts(alias string, password []byte) (
	key crypto.PrivateKey, certs []*x509.Certificate, err error) {

	entry := ks.entries[alias]
	if entry == nil {
		return
	}
	switch t := entry.(type) {
	case *privateKeyEntry:
		if len(t.certs) < 1 {
			return nil, nil, fmt.Errorf("key has no certificates")
		}
		key, err = t.Recover(password)
		if err == nil {
			certs = t.certs
			return
		}
	}
	return
}

// GetCert retrieves the specified certificate. Returns nil if the
// certificate does not exist or alias points to a non certificate
// entry.
func (ks *KeyStore) GetCert(alias string) (*x509.Certificate, error) {
	entry := ks.entries[alias]
	if entry == nil {
		return nil, nil
	}
	switch t := entry.(type) {
	case *trustedCertEntry:
		return t.cert, nil
	}
	return nil, nil
}

// ListPrivateKeys lists the names of the private keys stored in the key store.
func (ks *KeyStore) ListPrivateKeys() []string {
	var r []string
	for k, v := range ks.entries {
		if _, ok := v.(*privateKeyEntry); ok {
			r = append(r, k)
		}
	}
	return r
}

// ListCerts lists the names of the certs stored in the key store.
func (ks *KeyStore) ListCerts() []string {
	var r []string
	for k, v := range ks.entries {
		if _, ok := v.(*trustedCertEntry); ok {
			r = append(r, k)
		}
	}
	return r
}

func (ks *KeyStore) String() string {
	var buf bytes.Buffer
	for k, v := range ks.entries {
		fmt.Fprintf(&buf, "%s\n", k)
		fmt.Fprintf(&buf, "  %s\n", v)
	}
	return buf.String()
}

// LoadFromFile loads the key store from the specified file.
func LoadFromFile(filename string, password []byte) (*KeyStore, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return LoadFromReader(file, password)
}

// LoadFromReader loads the key store from the specified file.
func LoadFromReader(reader io.Reader, password []byte) (*KeyStore, error) {
	ks := &KeyStore{
		entries: make(map[string]interface{}),
	}
	err := ks.Parse(reader, password)
	if err != nil {
		return nil, err
	}
	return ks, err
}
<commit_msg>Fix a tiny typo in a comment.<commit_after>/*-
 * Copyright 2016 Square Inc.
 *
 * 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 jceks parses JCEKS (Java Cryptogaphy Extension Key Store)
// files and extracts keys and certificates. This module only implements
// a fraction of the JCEKS cryptographic protocols. In particular, it
// implements the SHA1 signature verification of the key store and the
// PBEWithMD5AndDES3CBC cipher for encrypting private keys.
package jceks

import (
	"bytes"
	"crypto"
	"crypto/sha1"
	"crypto/subtle"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/binary"
	"fmt"
	"hash"
	"io"
	"os"
	"time"
)

const (
	jceksMagic   = 0xcececece
	jceksVersion = 0x02
	jksMagic     = 0xfeedfeed
)

var (
	oidKeyProtector = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 42, 2, 17, 1, 1}
	oidPublicKeyRSA = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1}
	oidPublicKeyEC  = asn1.ObjectIdentifier{1, 2, 840, 10045, 2, 1}
)

type encryptedPrivateKeyInfo struct {
	Algo         pkix.AlgorithmIdentifier
	EncryptedKey []byte
}

type privateKeyInfo struct {
	Version    int
	Algo       pkix.AlgorithmIdentifier
	PrivateKey []byte
}

type privateKeyEntry struct {
	date       time.Time
	encodedKey []byte
	certs      []*x509.Certificate
}

// From https://golang.org/src/crypto/x509/sec1.go (also: see RFC 5915)
type ecPrivateKey struct {
	Version       int
	PrivateKey    []byte
	NamedCurveOID asn1.ObjectIdentifier `asn1:"optional,explicit,tag:0"`
	PublicKey     asn1.BitString        `asn1:"optional,explicit,tag:1"`
}

func (e *privateKeyEntry) String() string {
	return fmt.Sprintf("private-key: %s", e.date)
}

func (e *privateKeyEntry) Recover(password []byte) (crypto.PrivateKey, error) {
	var eKey encryptedPrivateKeyInfo
	_, err := asn1.Unmarshal(e.encodedKey, &eKey)
	if err != nil {
		return nil, err
	}
	var decryptedKey []byte
	switch {
	case eKey.Algo.Algorithm.Equal(oidPBEWithMD5AndDES3CBC):
		decryptedKey, err = recoverPBEWithMD5AndDES3CBC(eKey.Algo, eKey.EncryptedKey, password)
	case eKey.Algo.Algorithm.Equal(oidKeyProtector):
		// JavaSoft proprietary key-protection algorithm (used to protect
		// private keys in the keystore implementation that comes with JDK
		// 1.2). We shouldn't need this.
		fallthrough
	default:
		return nil, fmt.Errorf("unsupported encrypted-private-key algorithm: %v", eKey.Algo.Algorithm)
	}
	if err != nil {
		return nil, err
	}

	var pKey privateKeyInfo
	if _, err := asn1.Unmarshal(decryptedKey, &pKey); err != nil {
		return nil, err
	}
	if pKey.Algo.Algorithm.Equal(oidPublicKeyRSA) {
		return x509.ParsePKCS1PrivateKey(pKey.PrivateKey)
	}
	if pKey.Algo.Algorithm.Equal(oidPublicKeyEC) {
		// In a JCEKS file, the EC private key blob contains only the key itself, without the
		// named curve OID. Instead, the named curve OID is in a separate field (in the algorithm
		// indentifier for the keystore entry). However to parse the EC key properly we need the
		// EC key blob to have the curve OID...
		key := ecPrivateKey{}
		oid := asn1.ObjectIdentifier{}
		// Parse EC private key
		_, err := asn1.Unmarshal(pKey.PrivateKey, &key)
		if err != nil {
			return nil, fmt.Errorf("problem parsing ec key asn.1 struct: %s", err)
		}
		// Parse named curve OID from algorithm identifier
		_, err = asn1.Unmarshal(pKey.Algo.Parameters.FullBytes, &oid)
		if err != nil {
			return nil, fmt.Errorf("problem parsing ec key asn.1 struct: %s", err)
		}
		// Update key to add named curve info, re-marshal, and parse
		key.NamedCurveOID = oid
		raw, _ := asn1.Marshal(key)
		return x509.ParseECPrivateKey(raw)
	}
	return nil, fmt.Errorf("unsupported private-key algorithm: %v", pKey.Algo.Algorithm)
}

type trustedCertEntry struct {
	date time.Time
	cert *x509.Certificate
}

func (e *trustedCertEntry) String() string {
	return fmt.Sprintf("trusted-cert: %s", e.date)
}

// KeyStore maintains a map from alias name to the entry for that
// alias. Entries are currently either privateKeyEntry or
// trustedCertEntry.
type KeyStore struct {
	entries map[string]interface{}
}

// readUTF reads a java encoded UTF-8 string. The encoding provides a
// 2-byte prefix indicating the length of the string.
func readUTF(r io.Reader) (string, error) {
	var length uint16
	err := binary.Read(r, binary.BigEndian, &length)
	if err != nil {
		return "", err
	}
	buf := make([]byte, length)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return "", err
	}
	return string(buf), nil
}

// readBytes reads a byte array from the reader. The encoding provides
// a 4-byte prefix indicating the number of bytes which follow.
func readBytes(r io.Reader) ([]byte, error) {
	length, err := readInt32(r)
	if err != nil {
		return nil, err
	}
	buf := make([]byte, length)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	return buf, nil
}

func readInt32(r io.Reader) (int32, error) {
	var v int32
	err := binary.Read(r, binary.BigEndian, &v)
	return v, err
}

func readUint32(r io.Reader) (uint32, error) {
	var v uint32
	err := binary.Read(r, binary.BigEndian, &v)
	return v, err
}

func readDate(r io.Reader) (time.Time, error) {
	var v int64
	err := binary.Read(r, binary.BigEndian, &v)
	if err != nil {
		return time.Time{}, err
	}
	sec := v / 1000
	nsec := (v - sec*1000) * 1000 * 1000
	return time.Unix(sec, nsec), nil
}

// Returns a SHA1 hash which has been pre-keyed with the specified
// password according to the JCEKS algorithm.
func getPreKeyedHash(password []byte) hash.Hash {
	md := sha1.New()
	buf := make([]byte, len(password)*2)
	for i := 0; i < len(password); i++ {
		buf[i*2+1] = password[i]
	}
	md.Write(buf)
	// Yes, "Mighty Aphrodite" is a constant used by this method.
	md.Write([]byte("Mighty Aphrodite"))
	return md
}

func parseHeader(r io.Reader) (uint32, error) {
	magic, err := readUint32(r)
	if err != nil {
		return 0, err
	}
	if magic != jceksMagic && magic != jksMagic {
		return 0, fmt.Errorf("unexpected magic: %08x != (%08x || %08x)", magic, uint32(jceksMagic), uint32(jksMagic))
	}
	version, err := readUint32(r)
	if err != nil {
		return 0, err
	}
	return version, nil
}

func (ks *KeyStore) parsePrivateKey(r io.Reader) error {
	alias, err := readUTF(r)
	if err != nil {
		return err
	}
	entry := &privateKeyEntry{
		certs: []*x509.Certificate{},
	}
	entry.date, err = readDate(r)
	if err != nil {
		return err
	}
	entry.encodedKey, err = readBytes(r)
	if err != nil {
		return err
	}
	nCerts, err := readInt32(r)
	if err != nil {
		return err
	}
	for j := 0; j < int(nCerts); j++ {
		certType, err := readUTF(r)
		if err != nil {
			return err
		}
		if certType != "X.509" {
			return fmt.Errorf("unable to handle certificate type: %s", certType)
		}
		certBytes, err := readBytes(r)
		if err != nil {
			return err
		}
		cert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			return nil
		}
		entry.certs = append(entry.certs, cert)
	}
	ks.entries[alias] = entry
	return nil
}

func (ks *KeyStore) parseTrustedCert(r io.Reader) error {
	alias, err := readUTF(r)
	if err != nil {
		return err
	}
	entry := &trustedCertEntry{}
	entry.date, err = readDate(r)
	if err != nil {
		return err
	}

	certType, err := readUTF(r)
	if err != nil {
		return err
	}
	if certType != "X.509" {
		return fmt.Errorf("unable to handle certificate type: %s", certType)
	}
	certBytes, err := readBytes(r)
	if err != nil {
		return err
	}
	entry.cert, err = x509.ParseCertificate(certBytes)
	if err != nil {
		return nil
	}

	ks.entries[alias] = entry
	return nil
}

// Parse parses the key store from the specified reader.
func (ks *KeyStore) Parse(r io.Reader, password []byte) error {
	var md hash.Hash
	if password != nil {
		md = getPreKeyedHash(password)
		r = io.TeeReader(r, md)
	}

	version, err := parseHeader(r)
	if err != nil {
		return err
	}
	if version != jceksVersion {
		return fmt.Errorf("unexpected version: %d != %d", version, jceksVersion)
	}

	count, err := readInt32(r)
	if err != nil {
		return err
	}
	for i := 0; i < int(count); i++ {
		tag, err := readInt32(r)
		if err != nil {
			return err
		}
		switch tag {
		case 1:
			// Private-key entry
			err := ks.parsePrivateKey(r)
			if err != nil {
				return err
			}
		case 2:
			// Trusted-cert entry
			err := ks.parseTrustedCert(r)
			if err != nil {
				return err
			}
		case 3:
			// Secret-key entry
			return fmt.Errorf("unimplemented: Secret-key")
		default:
			return fmt.Errorf("unimplemented tag: %d", tag)
		}
	}

	if md != nil {
		computed := md.Sum([]byte{})
		actual := make([]byte, len(computed))
		_, err := io.ReadFull(r, actual)
		if err != nil {
			return err
		}
		if subtle.ConstantTimeCompare(computed, actual) != 1 {
			return fmt.Errorf("keystore was tampered with or password was incorrect")
		}
	}

	return nil
}

// GetPrivateKeyAndCerts retrieves the specified private key. Returns
// nil if the private key does not exist or alias points to a non
// private key entry.
func (ks *KeyStore) GetPrivateKeyAndCerts(alias string, password []byte) (
	key crypto.PrivateKey, certs []*x509.Certificate, err error) {

	entry := ks.entries[alias]
	if entry == nil {
		return
	}
	switch t := entry.(type) {
	case *privateKeyEntry:
		if len(t.certs) < 1 {
			return nil, nil, fmt.Errorf("key has no certificates")
		}
		key, err = t.Recover(password)
		if err == nil {
			certs = t.certs
			return
		}
	}
	return
}

// GetCert retrieves the specified certificate. Returns nil if the
// certificate does not exist or alias points to a non certificate
// entry.
func (ks *KeyStore) GetCert(alias string) (*x509.Certificate, error) {
	entry := ks.entries[alias]
	if entry == nil {
		return nil, nil
	}
	switch t := entry.(type) {
	case *trustedCertEntry:
		return t.cert, nil
	}
	return nil, nil
}

// ListPrivateKeys lists the names of the private keys stored in the key store.
func (ks *KeyStore) ListPrivateKeys() []string {
	var r []string
	for k, v := range ks.entries {
		if _, ok := v.(*privateKeyEntry); ok {
			r = append(r, k)
		}
	}
	return r
}

// ListCerts lists the names of the certs stored in the key store.
func (ks *KeyStore) ListCerts() []string {
	var r []string
	for k, v := range ks.entries {
		if _, ok := v.(*trustedCertEntry); ok {
			r = append(r, k)
		}
	}
	return r
}

func (ks *KeyStore) String() string {
	var buf bytes.Buffer
	for k, v := range ks.entries {
		fmt.Fprintf(&buf, "%s\n", k)
		fmt.Fprintf(&buf, "  %s\n", v)
	}
	return buf.String()
}

// LoadFromFile loads the key store from the specified file.
func LoadFromFile(filename string, password []byte) (*KeyStore, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return LoadFromReader(file, password)
}

// LoadFromReader loads the key store from the specified file.
func LoadFromReader(reader io.Reader, password []byte) (*KeyStore, error) {
	ks := &KeyStore{
		entries: make(map[string]interface{}),
	}
	err := ks.Parse(reader, password)
	if err != nil {
		return nil, err
	}
	return ks, err
}
<|endoftext|>


<commit_before>package quasar

import (
	"github.com/f483/dejavu"
	"math/rand"
	"sync"
	"time"
)

// FIXME use io.Reader and io.Writer where possible

type Quasar struct {
	net               networkOverlay
	subscribers       map[hash160digest][]chan []byte
	topics            map[hash160digest][]byte
	mutex             *sync.RWMutex
	peers             map[pubkey]*peerData
	history           dejavu.DejaVu // memory of past events
	cfg               config
	stopDispatcher    chan bool
	stopPropagation   chan bool
	stopExpiredPeerGC chan bool
}

// Create new quasar instance
func NewQuasar() *Quasar {
	return newQuasar(nil, defaultConfig) // FIXME add default network
}

func newQuasar(net networkOverlay, c config) *Quasar {
	d := dejavu.NewProbabilistic(c.historyLimit, c.historyAccuracy)
	return &Quasar{
		net:               net,
		subscribers:       make(map[hash160digest][]chan []byte),
		topics:            make(map[hash160digest][]byte),
		mutex:             new(sync.RWMutex),
		peers:             make(map[pubkey]*peerData),
		history:           d,
		cfg:               c,
		stopDispatcher:    nil, // set on Start() call
		stopPropagation:   nil, // set on Start() call
		stopExpiredPeerGC: nil, // set on Start() call
	}
}

func (q *Quasar) processUpdate(u *peerUpdate) {
	q.mutex.Lock()
	data, ok := q.peers[*u.peer]

	if !ok { // init if doesnt exist
		depth := q.cfg.filtersDepth
		data = &peerData{
			filters:    newFilters(q.cfg),
			timestamps: make([]uint64, depth, depth),
		}
		q.peers[*u.peer] = data
	}

	// update peer data
	data.filters[u.index] = u.filter
	data.timestamps[u.index] = uint64(time.Now().Unix())
	q.mutex.Unlock()
}

func (q *Quasar) Publish(topic []byte, message []byte) {
	// TODO validate input
	go q.route(newEvent(topic, message, q.cfg.defaultEventTTL))
}

func (q *Quasar) isDuplicate(e *event) bool {
	return q.history.Witness(append(e.topicDigest[:20], e.message...))
}

func (q *Quasar) deliver(receivers []chan []byte, e *event) {
	for _, receiver := range receivers {
		receiver <- e.message
	}
}

// Algorithm 1 from the quasar paper.
func (q *Quasar) sendUpdates() {
	q.mutex.RLock()
	filters := newFilters(q.cfg)
	for digest := range q.subscribers {
		filters[0] = filterInsertDigest(filters[0], digest, q.cfg)
	}
	pubkey := q.net.Id()
	filters[0] = filterInsert(filters[0], pubkey[:], q.cfg)
	for _, data := range q.peers {
		if peerDataExpired(data, q.cfg) {
			continue
		}
		for i := 1; uint32(i) < q.cfg.filtersDepth; i++ {
			filters[i] = mergeFilters(filters[i], data.filters[i-1])
		}
	}
	for _, peerId := range q.net.ConnectedPeers() {
		for i := 0; uint32(i) < (q.cfg.filtersDepth - 1); i++ {
			// top filter never sent as not used by peers
			go q.net.SendUpdate(&peerId, uint32(i), filters[i])
		}
	}
	q.mutex.RUnlock()
}

// Algorithm 2 from the quasar paper.
func (q *Quasar) route(e *event) {
	q.mutex.RLock()
	if q.isDuplicate(e) {
		q.mutex.RUnlock()
		return
	}
	if receivers, ok := q.subscribers[*e.topicDigest]; ok {
		q.deliver(receivers, e)
		e.publishers = append(e.publishers, q.net.Id())
		for _, peerId := range q.net.ConnectedPeers() {
			go q.net.SendEvent(&peerId, e)
		}
		q.mutex.RUnlock()
		return
	}
	e.ttl -= 1
	if e.ttl == 0 {
		q.mutex.RUnlock()
		return
	}
	for i := 0; uint32(i) < q.cfg.filtersDepth; i++ {
		for peerId, data := range q.peers {
			f := data.filters[i]
			if filterContainsDigest(f, *e.topicDigest, q.cfg) {
				negRt := false
				for _, publisher := range e.publishers {
					if filterContains(f, publisher[:], q.cfg) {
						negRt = true
					}
				}
				if !negRt {
					go q.net.SendEvent(&peerId, e)
					q.mutex.RUnlock()
					return
				}
			}
		}
	}
	peerId := q.randomPeer()
	if peerId != nil {
		go q.net.SendEvent(peerId, e)
	}
	q.mutex.RUnlock()
}

func (q *Quasar) randomPeer() *pubkey {
	peers := q.net.ConnectedPeers()
	if len(peers) == 0 {
		return nil
	}
	peerId := peers[rand.Intn(len(peers))]
	return &peerId
}

func (q *Quasar) dispatchInput() {
	for {
		select {
		case peerUpdate := <-q.net.ReceivedUpdateChannel():
			if validUpdate(peerUpdate, q.cfg) {
				go q.processUpdate(peerUpdate)
			}
		case event := <-q.net.ReceivedEventChannel():
			if validEvent(event) {
				go q.route(event)
			}
		case <-q.stopDispatcher:
			return
		}
	}
}

func (q *Quasar) removeExpiredPeers() {
	q.mutex.Lock()
	toRemove := []*pubkey{}
	for peerId, data := range q.peers {
		if peerDataExpired(data, q.cfg) {
			toRemove = append(toRemove, &peerId)
		}
	}
	for _, peerId := range toRemove {
		delete(q.peers, *peerId)
	}
	q.mutex.Unlock()
}

func (q *Quasar) expiredPeerGC() {
	delay := time.Duration(1) * time.Millisecond
	for {
		select {
		case <-time.After(delay):
			go q.removeExpiredPeers()
		case <-q.stopExpiredPeerGC:
			return
		}
	}
}

func (q *Quasar) propagateFilters() {
	delay := time.Duration(q.cfg.propagationDelay) * time.Second
	for {
		select {
		case <-time.After(delay):
			go q.sendUpdates()
		case <-q.stopPropagation:
			return
		}
	}
}

// Start quasar system
func (q *Quasar) Start() {
	q.net.Start()
	q.stopDispatcher = make(chan bool)
	q.stopPropagation = make(chan bool)
	q.stopExpiredPeerGC = make(chan bool)
	go q.dispatchInput()
	go q.propagateFilters()
	go q.expiredPeerGC()
}

// Stop quasar system
func (q *Quasar) Stop() {
	q.net.Stop()
	q.stopDispatcher <- true
	q.stopPropagation <- true
	q.stopExpiredPeerGC <- true
}

// Subscribe provided message receiver channel to given topic.
func (q *Quasar) Subscribe(topic []byte, receiver chan []byte) {
	// TODO validate input
	digest := hash160(topic)
	q.mutex.Lock()
	receivers, ok := q.subscribers[digest]
	if ok != true { // new subscription
		q.subscribers[digest] = []chan []byte{receiver}
		q.topics[digest] = topic
	} else { // append to existing subscribers
		q.subscribers[digest] = append(receivers, receiver)
	}
	q.mutex.Unlock()
}

// Unsubscribe message receiver channel from topic. If nil receiver
// channel is provided all message receiver channels for given topic
// will be removed.
func (q *Quasar) Unsubscribe(topic []byte, receiver chan []byte) {
	// TODO validate input

	digest := hash160(topic)
	q.mutex.Lock()
	receivers, ok := q.subscribers[digest]

	// remove specific message receiver
	if ok && receiver != nil {
		for i, v := range receivers {
			if v == receiver {
				receivers = append(receivers[:i], receivers[i+1:]...)
				q.subscribers[digest] = receivers
				break
			}
		}
	}

	// remove sub key if no specific message
	// receiver provided or no message receiver remaining
	if ok && (receiver == nil || len(q.subscribers[digest]) == 0) {
		delete(q.subscribers, digest)
		delete(q.topics, digest)
	}
	q.mutex.Unlock()
}

// Subscribers retruns message receivers for given topic.
func (q *Quasar) Subscribers(topic []byte) []chan []byte {
	// TODO validate input
	digest := hash160(topic)
	results := []chan []byte{}
	q.mutex.RLock()
	if receivers, ok := q.subscribers[digest]; ok {
		results = append(results, receivers...)
	}
	q.mutex.RUnlock()
	return results
}

// SubscribedTopics retruns a slice of currently subscribed topics.
func (q *Quasar) Subscriptions() [][]byte {
	q.mutex.RLock()
	topics := make([][]byte, len(q.topics))
	i := 0
	for _, topic := range q.topics {
		topics[i] = topic
		i++
	}
	q.mutex.RUnlock()
	return topics
}
<commit_msg>added missing documentation<commit_after>package quasar

import (
	"github.com/f483/dejavu"
	"math/rand"
	"sync"
	"time"
)

// FIXME use io.Reader and io.Writer where possible

// Quasar holds the pubsup state
type Quasar struct {
	net               networkOverlay
	subscribers       map[hash160digest][]chan []byte
	topics            map[hash160digest][]byte
	mutex             *sync.RWMutex
	peers             map[pubkey]*peerData
	history           dejavu.DejaVu // memory of past events
	cfg               config
	stopDispatcher    chan bool
	stopPropagation   chan bool
	stopExpiredPeerGC chan bool
}

// NewQuasar create new instance with the sane defaults.
func NewQuasar() *Quasar {
	// FIXME enable passing of nodeId/pubkey
	return newQuasar(nil, defaultConfig) // FIXME add default network
}

func newQuasar(net networkOverlay, c config) *Quasar {
	d := dejavu.NewProbabilistic(c.historyLimit, c.historyAccuracy)
	return &Quasar{
		net:               net,
		subscribers:       make(map[hash160digest][]chan []byte),
		topics:            make(map[hash160digest][]byte),
		mutex:             new(sync.RWMutex),
		peers:             make(map[pubkey]*peerData),
		history:           d,
		cfg:               c,
		stopDispatcher:    nil, // set on Start() call
		stopPropagation:   nil, // set on Start() call
		stopExpiredPeerGC: nil, // set on Start() call
	}
}

func (q *Quasar) processUpdate(u *peerUpdate) {
	q.mutex.Lock()
	data, ok := q.peers[*u.peer]

	if !ok { // init if doesnt exist
		depth := q.cfg.filtersDepth
		data = &peerData{
			filters:    newFilters(q.cfg),
			timestamps: make([]uint64, depth, depth),
		}
		q.peers[*u.peer] = data
	}

	// update peer data
	data.filters[u.index] = u.filter
	data.timestamps[u.index] = uint64(time.Now().Unix())
	q.mutex.Unlock()
}

// Publish a message on the network for given topic.
func (q *Quasar) Publish(topic []byte, message []byte) {
	// TODO validate input
	go q.route(newEvent(topic, message, q.cfg.defaultEventTTL))
}

func (q *Quasar) isDuplicate(e *event) bool {
	return q.history.Witness(append(e.topicDigest[:20], e.message...))
}

func (q *Quasar) deliver(receivers []chan []byte, e *event) {
	for _, receiver := range receivers {
		receiver <- e.message
	}
}

// Algorithm 1 from the quasar paper.
func (q *Quasar) sendUpdates() {
	q.mutex.RLock()
	filters := newFilters(q.cfg)
	for digest := range q.subscribers {
		filters[0] = filterInsertDigest(filters[0], digest, q.cfg)
	}
	pubkey := q.net.Id()
	filters[0] = filterInsert(filters[0], pubkey[:], q.cfg)
	for _, data := range q.peers {
		if peerDataExpired(data, q.cfg) {
			continue
		}
		for i := 1; uint32(i) < q.cfg.filtersDepth; i++ {
			filters[i] = mergeFilters(filters[i], data.filters[i-1])
		}
	}
	for _, peerId := range q.net.ConnectedPeers() {
		for i := 0; uint32(i) < (q.cfg.filtersDepth - 1); i++ {
			// top filter never sent as not used by peers
			go q.net.SendUpdate(&peerId, uint32(i), filters[i])
		}
	}
	q.mutex.RUnlock()
}

// Algorithm 2 from the quasar paper.
func (q *Quasar) route(e *event) {
	q.mutex.RLock()
	if q.isDuplicate(e) {
		q.mutex.RUnlock()
		return
	}
	if receivers, ok := q.subscribers[*e.topicDigest]; ok {
		q.deliver(receivers, e)
		e.publishers = append(e.publishers, q.net.Id())
		for _, peerId := range q.net.ConnectedPeers() {
			go q.net.SendEvent(&peerId, e)
		}
		q.mutex.RUnlock()
		return
	}
	e.ttl -= 1
	if e.ttl == 0 {
		q.mutex.RUnlock()
		return
	}
	for i := 0; uint32(i) < q.cfg.filtersDepth; i++ {
		for peerId, data := range q.peers {
			f := data.filters[i]
			if filterContainsDigest(f, *e.topicDigest, q.cfg) {
				negRt := false
				for _, publisher := range e.publishers {
					if filterContains(f, publisher[:], q.cfg) {
						negRt = true
					}
				}
				if !negRt {
					go q.net.SendEvent(&peerId, e)
					q.mutex.RUnlock()
					return
				}
			}
		}
	}
	peerId := q.randomPeer()
	if peerId != nil {
		go q.net.SendEvent(peerId, e)
	}
	q.mutex.RUnlock()
}

func (q *Quasar) randomPeer() *pubkey {
	peers := q.net.ConnectedPeers()
	if len(peers) == 0 {
		return nil
	}
	peerId := peers[rand.Intn(len(peers))]
	return &peerId
}

func (q *Quasar) dispatchInput() {
	for {
		select {
		case peerUpdate := <-q.net.ReceivedUpdateChannel():
			if validUpdate(peerUpdate, q.cfg) {
				go q.processUpdate(peerUpdate)
			}
		case event := <-q.net.ReceivedEventChannel():
			if validEvent(event) {
				go q.route(event)
			}
		case <-q.stopDispatcher:
			return
		}
	}
}

func (q *Quasar) removeExpiredPeers() {
	q.mutex.Lock()
	toRemove := []*pubkey{}
	for peerId, data := range q.peers {
		if peerDataExpired(data, q.cfg) {
			toRemove = append(toRemove, &peerId)
		}
	}
	for _, peerId := range toRemove {
		delete(q.peers, *peerId)
	}
	q.mutex.Unlock()
}

func (q *Quasar) expiredPeerGC() {
	delay := time.Duration(1) * time.Millisecond
	for {
		select {
		case <-time.After(delay):
			go q.removeExpiredPeers()
		case <-q.stopExpiredPeerGC:
			return
		}
	}
}

func (q *Quasar) propagateFilters() {
	delay := time.Duration(q.cfg.propagationDelay) * time.Second
	for {
		select {
		case <-time.After(delay):
			go q.sendUpdates()
		case <-q.stopPropagation:
			return
		}
	}
}

// Start quasar system
func (q *Quasar) Start() {
	q.net.Start()
	q.stopDispatcher = make(chan bool)
	q.stopPropagation = make(chan bool)
	q.stopExpiredPeerGC = make(chan bool)
	go q.dispatchInput()
	go q.propagateFilters()
	go q.expiredPeerGC()
}

// Stop quasar system
func (q *Quasar) Stop() {
	q.net.Stop()
	q.stopDispatcher <- true
	q.stopPropagation <- true
	q.stopExpiredPeerGC <- true
}

// Subscribe provided message receiver channel to given topic.
func (q *Quasar) Subscribe(topic []byte, receiver chan []byte) {
	// TODO validate input
	digest := hash160(topic)
	q.mutex.Lock()
	receivers, ok := q.subscribers[digest]
	if ok != true { // new subscription
		q.subscribers[digest] = []chan []byte{receiver}
		q.topics[digest] = topic
	} else { // append to existing subscribers
		q.subscribers[digest] = append(receivers, receiver)
	}
	q.mutex.Unlock()
}

// Unsubscribe message receiver channel from topic. If nil receiver
// channel is provided all message receiver channels for given topic
// will be removed.
func (q *Quasar) Unsubscribe(topic []byte, receiver chan []byte) {
	// TODO validate input

	digest := hash160(topic)
	q.mutex.Lock()
	receivers, ok := q.subscribers[digest]

	// remove specific message receiver
	if ok && receiver != nil {
		for i, v := range receivers {
			if v == receiver {
				receivers = append(receivers[:i], receivers[i+1:]...)
				q.subscribers[digest] = receivers
				break
			}
		}
	}

	// remove sub key if no specific message
	// receiver provided or no message receiver remaining
	if ok && (receiver == nil || len(q.subscribers[digest]) == 0) {
		delete(q.subscribers, digest)
		delete(q.topics, digest)
	}
	q.mutex.Unlock()
}

// Subscribers retruns message receivers for given topic.
func (q *Quasar) Subscribers(topic []byte) []chan []byte {
	// TODO validate input
	digest := hash160(topic)
	results := []chan []byte{}
	q.mutex.RLock()
	if receivers, ok := q.subscribers[digest]; ok {
		results = append(results, receivers...)
	}
	q.mutex.RUnlock()
	return results
}

// Subscriptions retruns a slice of currently subscribed topics.
func (q *Quasar) Subscriptions() [][]byte {
	q.mutex.RLock()
	topics := make([][]byte, len(q.topics))
	i := 0
	for _, topic := range q.topics {
		topics[i] = topic
		i++
	}
	q.mutex.RUnlock()
	return topics
}
<|endoftext|>


<commit_before>package incognitomail

import (
	"errors"
	"time"

	"github.com/boltdb/bolt"
)

// IncognitoData holds a "connection" to the persistence layer. To create a valid IncognitoData object, call OpenIncognitoData().
type IncognitoData struct {
	db *bolt.DB
}

const (
	targetsBucketName  = "targets"
	accountsBucketName = "accounts"
	handlesBucketName  = "handles"
)

var (
	// ErrAccountNotFound is used when an action requires an account to exist, but it wasn't found.
	ErrAccountNotFound = errors.New("account not found")

	// ErrAccountExists is used when trying to create an account with a given secret, but it already exists.
	ErrAccountExists = errors.New("account already exists")

	// ErrHandleNotFound is used when an action requires a handle to exist, but it wasn't found.
	ErrHandleNotFound = errors.New("handle not found")

	// ErrHandleExists is used when trying to create a handle, but it already exists.
	ErrHandleExists = errors.New("handle already exists")
)

// OpenIncognitoData returns an IncognitoData object with a successful "connection" to the persistence layer, ready to be used.
func OpenIncognitoData() (*IncognitoData, error) {
	db, err := bolt.Open(Config.Persistence.DatabasePath, 0600, nil)

	if err != nil {
		return nil, err
	}

	// Create "static" buckets that are used for persistence
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(targetsBucketName))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte(accountsBucketName))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte(handlesBucketName))
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return &IncognitoData{
		db: db,
	}, nil
}

// NewAccount generates a new account with the given secret and target email address.
func (a *IncognitoData) NewAccount(secret, target string) error {
	err := a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))

		if b != nil {
			return ErrAccountExists
		}

		_, err := tx.CreateBucket([]byte(secret))
		if err != nil {
			return err
		}

		b = tx.Bucket([]byte(targetsBucketName))
		err = b.Put([]byte(secret), []byte(target))
		if err != nil {
			return err
		}

		now, err := time.Now().GobEncode()

		if err != nil {
			return err
		}

		b = tx.Bucket([]byte(accountsBucketName))
		err = b.Put([]byte(secret), now)
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return err
	}

	return nil
}

// DeleteAccount deletes all information related to the account with the given secret. If no account with that secret exists, it does nothing.
func (a *IncognitoData) DeleteAccount(secret string) {
	a.db.Update(func(tx *bolt.Tx) error {
		tx.DeleteBucket([]byte(secret))

		b := tx.Bucket([]byte(targetsBucketName))
		b.Delete([]byte(secret))

		b = tx.Bucket([]byte(accountsBucketName))
		b.Delete([]byte(secret))

		return nil
	})
}

// NewAccountHandle stores the given handle for the account with the given secret.
func (a *IncognitoData) NewAccountHandle(secret, handle string) error {
	err := a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))
		if b == nil {
			return ErrAccountNotFound
		}

		hb := tx.Bucket([]byte(handlesBucketName))
		h := hb.Get([]byte(handle))
		if h != nil {
			return ErrHandleExists
		}

		now, err := time.Now().GobEncode()
		if err != nil {
			return err
		}

		err = b.Put([]byte(handle), now)
		if err != nil {
			return err
		}

		err = hb.Put([]byte(handle), now)
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return err
	}

	return nil
}

// DeleteAccountHandle deletes the given handle from the account with the given secret. If either the account or the handle does not exist, this does nothing.
func (a *IncognitoData) DeleteAccountHandle(secret, handle string) {
	// Note that we still return errors from the following func, but don't care about them
	a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))
		if b == nil {
			return ErrAccountNotFound
		}

		b.Delete([]byte(handle))
		return nil
	})
}

// GetAccountTarget returns the target registered for the account with the given secret.
func (a *IncognitoData) GetAccountTarget(secret string) (string, error) {
	var target string

	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(targetsBucketName))
		t := b.Get([]byte(secret))
		if t == nil {
			return ErrAccountNotFound
		}

		// Note: boltdb only keeps the value of t until the transaction ends, so we must copy it somewhere else now.
		// However, the call to string(t) internally does that for us, as it will ultimately call copy() to copy the values to a new byte slice for the resulting string.
		target = string(t)
		return nil
	})

	if err != nil {
		return "", err
	}

	return target, nil
}

// HasAccount returns true if an account with the given secret exists, false otherwise.
func (a *IncognitoData) HasAccount(secret string) bool {
	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(accountsBucketName))
		t := b.Get([]byte(secret))
		if t == nil {
			return ErrAccountNotFound
		}

		return nil
	})

	return err == nil
}

// HasHandleGlobal returns true if the given handle exists for any account, false otherwise.
func (a *IncognitoData) HasHandleGlobal(handle string) bool {
	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(handlesBucketName))
		t := b.Get([]byte(handle))
		if t == nil {
			return ErrAccountNotFound
		}

		return nil
	})

	return err == nil
}

// ListAccountHandles returns an array with all handles from the account with the given secret.
func (a *IncognitoData) ListAccountHandles(secret string) ([]string, error) {
	var result []string

	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))

		b.ForEach(func(k, v []byte) error {
			// Note: boltdb only keeps the values of k and v until the transaction ends, so we must copy these values somewhere else now.
			// However, the call to string(k) internally does that for us, as it will ultimately call copy() to copy the values to a new byte slice for the resulting string.
			result = append(result, string(k))
			return nil
		})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return result, nil
}

// Close closes the "connection" with the persistence layer.
func (a *IncognitoData) Close() {
	a.db.Close()
}
<commit_msg>Fix bug when deleting handles.<commit_after>package incognitomail

import (
	"errors"
	"time"

	"github.com/boltdb/bolt"
)

// IncognitoData holds a "connection" to the persistence layer. To create a valid IncognitoData object, call OpenIncognitoData().
type IncognitoData struct {
	db *bolt.DB
}

const (
	targetsBucketName  = "targets"
	accountsBucketName = "accounts"
	handlesBucketName  = "handles"
)

var (
	// ErrAccountNotFound is used when an action requires an account to exist, but it wasn't found.
	ErrAccountNotFound = errors.New("account not found")

	// ErrAccountExists is used when trying to create an account with a given secret, but it already exists.
	ErrAccountExists = errors.New("account already exists")

	// ErrHandleNotFound is used when an action requires a handle to exist, but it wasn't found.
	ErrHandleNotFound = errors.New("handle not found")

	// ErrHandleExists is used when trying to create a handle, but it already exists.
	ErrHandleExists = errors.New("handle already exists")
)

// OpenIncognitoData returns an IncognitoData object with a successful "connection" to the persistence layer, ready to be used.
func OpenIncognitoData() (*IncognitoData, error) {
	db, err := bolt.Open(Config.Persistence.DatabasePath, 0600, nil)

	if err != nil {
		return nil, err
	}

	// Create "static" buckets that are used for persistence
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(targetsBucketName))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte(accountsBucketName))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte(handlesBucketName))
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return &IncognitoData{
		db: db,
	}, nil
}

// NewAccount generates a new account with the given secret and target email address.
func (a *IncognitoData) NewAccount(secret, target string) error {
	err := a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))

		if b != nil {
			return ErrAccountExists
		}

		_, err := tx.CreateBucket([]byte(secret))
		if err != nil {
			return err
		}

		b = tx.Bucket([]byte(targetsBucketName))
		err = b.Put([]byte(secret), []byte(target))
		if err != nil {
			return err
		}

		now, err := time.Now().GobEncode()

		if err != nil {
			return err
		}

		b = tx.Bucket([]byte(accountsBucketName))
		err = b.Put([]byte(secret), now)
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return err
	}

	return nil
}

// DeleteAccount deletes all information related to the account with the given secret. If no account with that secret exists, it does nothing.
func (a *IncognitoData) DeleteAccount(secret string) {
	a.db.Update(func(tx *bolt.Tx) error {
		tx.DeleteBucket([]byte(secret))

		b := tx.Bucket([]byte(targetsBucketName))
		b.Delete([]byte(secret))

		b = tx.Bucket([]byte(accountsBucketName))
		b.Delete([]byte(secret))

		return nil
	})
}

// NewAccountHandle stores the given handle for the account with the given secret.
func (a *IncognitoData) NewAccountHandle(secret, handle string) error {
	err := a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))
		if b == nil {
			return ErrAccountNotFound
		}

		hb := tx.Bucket([]byte(handlesBucketName))
		h := hb.Get([]byte(handle))
		if h != nil {
			return ErrHandleExists
		}

		now, err := time.Now().GobEncode()
		if err != nil {
			return err
		}

		err = b.Put([]byte(handle), now)
		if err != nil {
			return err
		}

		err = hb.Put([]byte(handle), now)
		if err != nil {
			return err
		}

		return nil
	})

	if err != nil {
		return err
	}

	return nil
}

// DeleteAccountHandle deletes the given handle from the account with the given secret. If either the account or the handle does not exist, this does nothing.
func (a *IncognitoData) DeleteAccountHandle(secret, handle string) {
	// Note that we still return errors from the following func, but don't care about them
	a.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))
		if b == nil {
			return ErrAccountNotFound
		}

		// Also delete from the global handles name
		hb := tx.Bucket([]byte(handlesBucketName))

		b.Delete([]byte(handle))
		hb.Delete([]byte(handle))
		return nil
	})
}

// GetAccountTarget returns the target registered for the account with the given secret.
func (a *IncognitoData) GetAccountTarget(secret string) (string, error) {
	var target string

	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(targetsBucketName))
		t := b.Get([]byte(secret))
		if t == nil {
			return ErrAccountNotFound
		}

		// Note: boltdb only keeps the value of t until the transaction ends, so we must copy it somewhere else now.
		// However, the call to string(t) internally does that for us, as it will ultimately call copy() to copy the values to a new byte slice for the resulting string.
		target = string(t)
		return nil
	})

	if err != nil {
		return "", err
	}

	return target, nil
}

// HasAccount returns true if an account with the given secret exists, false otherwise.
func (a *IncognitoData) HasAccount(secret string) bool {
	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(accountsBucketName))
		t := b.Get([]byte(secret))
		if t == nil {
			return ErrAccountNotFound
		}

		return nil
	})

	return err == nil
}

// HasHandleGlobal returns true if the given handle exists for any account, false otherwise.
func (a *IncognitoData) HasHandleGlobal(handle string) bool {
	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(handlesBucketName))
		t := b.Get([]byte(handle))
		if t == nil {
			return ErrAccountNotFound
		}

		return nil
	})

	return err == nil
}

// ListAccountHandles returns an array with all handles from the account with the given secret.
func (a *IncognitoData) ListAccountHandles(secret string) ([]string, error) {
	var result []string

	err := a.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(secret))

		b.ForEach(func(k, v []byte) error {
			// Note: boltdb only keeps the values of k and v until the transaction ends, so we must copy these values somewhere else now.
			// However, the call to string(k) internally does that for us, as it will ultimately call copy() to copy the values to a new byte slice for the resulting string.
			result = append(result, string(k))
			return nil
		})

		return nil
	})

	if err != nil {
		return nil, err
	}

	return result, nil
}

// Close closes the "connection" with the persistence layer.
func (a *IncognitoData) Close() {
	a.db.Close()
}
<|endoftext|>


<commit_before>package person

import (
	"crypto/md5"
	"database/sql/driver"
	"fmt"
	"io"

	"github.com/sourcegraph/go-nnz/nnz"
	"sourcegraph.com/sourcegraph/srcgraph/db_common"
)

// User represents a user.
type User struct {
	// UID is the numeric primary key for a user.
	UID UID `db:"uid"`

	// GitHubID is the numeric ID of the GitHub user account corresponding to
	// this user.
	GitHubID nnz.Int `db:"github_id"`

	// Login is the user's username, which typically corresponds to the user's
	// GitHub login.
	Login string

	// Name is the (possibly empty) full name of the user.
	Name string

	// AvatarURL is the URL to an avatar image specified by the user.
	AvatarURL string

	// Location is the user's physical location (from their GitHub profile).
	Location string `json:",omitempty"`

	// Company is the user's company (from their GitHub profile).
	Company string `json:",omitempty"`

	// HomepageURL is the user's homepage or blog URL (from their GitHub
	// profile).
	HomepageURL string `db:"homepage_url" json:",omitempty"`

	// Transient is if this user was constructed on the fly and is not persisted
	// or resolved to a Sourcegraph/GitHub/etc. user.
	Transient bool `db:"-" json:",omitempty"`

	// UserProfileDisabled is whether the user profile should not be displayed
	// on the Web app.
	UserProfileDisabled bool `db:"user_profile_disabled" json:",omitempty"`

	// RegisteredAt is the date that the user registered. If the user has not
	// registered (i.e., we have processed their repos but they haven't signed
	// into Sourcegraph), it is null.
	RegisteredAt db_common.NullTime `db:"registered_at"`

	// GitHubOAuth2AccessToken is the user's GitHub access token.
	GitHubOAuth2AccessToken string `db:"github_oauth2_access_token" json:"-"` // don't write the OAuth2 access token to JSON
}

// GitHubLogin returns the user's Login. They are the same for now, but callers
// that intend to get the GitHub login should call GitHubLogin() so that we can
// decouple the logins in the future if needed.
func (u *User) GitHubLogin() string {
	return u.Login
}

func (u *User) AvatarURLOfSize(width int) string {
	return u.AvatarURL + fmt.Sprintf("&s=%d", width)
}

// CanSync is whether this person can be synced with the external source that
// the person was originally fetched from (e.g., GitHub users).
func (u *User) CanSync() bool {
	return !u.Transient
}

// UID is the numeric primary key for a user.
type UID int

// Scan implements database/sql.Scanner.
func (x *UID) Scan(v interface{}) error {
	if data, ok := v.(int64); ok {
		*x = UID(data)
		return nil
	}
	return fmt.Errorf("%T.Scan failed: %v", x, v)
}

// Value implements database/sql/driver.Valuer.
func (x UID) Value() (driver.Value, error) {
	return int64(x), nil
}

// DefaultAvatarSize is the size, in pixels, of avatar images if no size is
// specified.
const DefaultAvatarSize = 128

// GravatarURL returns the URL to the Gravatar avatar image for email. If size
// is 0, DefaultAvatarSize is used.
func GravatarURL(email string, size uint16) string {
	if size == 0 {
		size = DefaultAvatarSize
	}
	h := md5.New()
	io.WriteString(h, email)
	return fmt.Sprintf("https://secure.gravatar.com/avatar/%x?s=%d&d=mm", h.Sum(nil), size)
}

// UserEmail is a row in the user_email DB table.
type UserEmail struct {
	UID   UID
	Email string
}
<commit_msg>repo contributions profile page<commit_after>package person

import (
	"crypto/md5"
	"database/sql/driver"
	"fmt"
	"io"

	"github.com/sourcegraph/go-nnz/nnz"
	"sourcegraph.com/sourcegraph/srcgraph/db_common"
)

// User represents a user.
type User struct {
	// UID is the numeric primary key for a user.
	UID UID `db:"uid"`

	// GitHubID is the numeric ID of the GitHub user account corresponding to
	// this user.
	GitHubID nnz.Int `db:"github_id"`

	// Login is the user's username, which typically corresponds to the user's
	// GitHub login.
	Login string

	// Name is the (possibly empty) full name of the user.
	Name string

	// AvatarURL is the URL to an avatar image specified by the user.
	AvatarURL string

	// Location is the user's physical location (from their GitHub profile).
	Location string `json:",omitempty"`

	// Company is the user's company (from their GitHub profile).
	Company string `json:",omitempty"`

	// HomepageURL is the user's homepage or blog URL (from their GitHub
	// profile).
	HomepageURL string `db:"homepage_url" json:",omitempty"`

	// Transient is if this user was constructed on the fly and is not persisted
	// or resolved to a Sourcegraph/GitHub/etc. user.
	Transient bool `db:"-" json:",omitempty"`

	// UserProfileDisabled is whether the user profile should not be displayed
	// on the Web app.
	UserProfileDisabled bool `db:"user_profile_disabled" json:",omitempty"`

	// RegisteredAt is the date that the user registered. If the user has not
	// registered (i.e., we have processed their repos but they haven't signed
	// into Sourcegraph), it is null.
	RegisteredAt db_common.NullTime `db:"registered_at"`

	// GitHubOAuth2AccessToken is the user's GitHub access token.
	GitHubOAuth2AccessToken string `db:"github_oauth2_access_token" json:"-"` // don't write the OAuth2 access token to JSON
}

// GitHubLogin returns the user's Login. They are the same for now, but callers
// that intend to get the GitHub login should call GitHubLogin() so that we can
// decouple the logins in the future if needed.
func (u *User) GitHubLogin() string {
	return u.Login
}

func (u *User) AvatarURLOfSize(width int) string {
	return u.AvatarURL + fmt.Sprintf("&s=%d", width)
}

// CanOwnRepositories is whether the user is capable of owning repositories
// (e.g., GitHub users can own GitHub repositories).
func (u *User) CanOwnRepositories() bool {
	return u.GitHubLogin() != ""
}

// CanSync is whether this person can be synced with the external source that
// the person was originally fetched from (e.g., GitHub users).
func (u *User) CanSync() bool {
	return !u.Transient
}

// UID is the numeric primary key for a user.
type UID int

// Scan implements database/sql.Scanner.
func (x *UID) Scan(v interface{}) error {
	if data, ok := v.(int64); ok {
		*x = UID(data)
		return nil
	}
	return fmt.Errorf("%T.Scan failed: %v", x, v)
}

// Value implements database/sql/driver.Valuer.
func (x UID) Value() (driver.Value, error) {
	return int64(x), nil
}

// DefaultAvatarSize is the size, in pixels, of avatar images if no size is
// specified.
const DefaultAvatarSize = 128

// GravatarURL returns the URL to the Gravatar avatar image for email. If size
// is 0, DefaultAvatarSize is used.
func GravatarURL(email string, size uint16) string {
	if size == 0 {
		size = DefaultAvatarSize
	}
	h := md5.New()
	io.WriteString(h, email)
	return fmt.Sprintf("https://secure.gravatar.com/avatar/%x?s=%d&d=mm", h.Sum(nil), size)
}

// UserEmail is a row in the user_email DB table.
type UserEmail struct {
	UID   UID
	Email string
}
<|endoftext|>


<commit_before>package gocube

import "errors"

// inverseXCornerIndices is the inverse permutation of xCornerIndices.
var inverseXCornerIndices []int = []int{2, 0, 3, 1, 6, 4, 7, 5}

// inverseZCornerIndices is the inverse permutation of zCornerIndices.
var inverseZCornerIndices []int = []int{4, 5, 0, 1, 6, 7, 2, 3}

// A Phase2Cube represents the parts of a cube that are important for phase-2
// solving.
type Phase2Cube struct {
	// CornerPermutation represents the permutation of the corners.
	CornerPermutation int

	// EdgePermutation represents the permutation of the 8 top/bottom edges.
	EdgePermutation int

	// SlicePermutation represents the permutation of the
	SlicePermutation int
}

// NewPhase2Cube generates a Phase2Cube from a CubieCube.
// The axis argument is 0 for X axis, 1 for Y axis, or 2 for Z axis.
// If the cube is not reduced to phase-2 in the given axis, this may return an
// error, but it does not validate everything.
func NewPhase2Cube(c CubieCube, axis int) (Phase2Cube, error) {
	var res Phase2Cube
	if axis == 0 {
		res.CornerPermutation = encodeXCornerPerm(&c.Corners)
		res.EdgePermutation = encodeRLEdges(&c.Edges)
		res.SlicePermutation = encodeMSlicePerm(&c.Edges)
	} else if axis == 1 {
		res.CornerPermutation = encodeYCornerPerm(&c.Corners)
		res.EdgePermutation = encodeUDEdges(&c.Edges)
		res.SlicePermutation = encodeESlicePerm(&c.Edges)
	} else {
		res.CornerPermutation = encodeZCornerPerm(&c.Corners)
		res.EdgePermutation = encodeFBEdges(&c.Edges)
		res.SlicePermutation = encodeSSlicePerm(&c.Edges)
	}
	if res.EdgePermutation < 0 {
		return res, errors.New("invalid edge permutation")
	} else if res.SlicePermutation < 0 {
		return res, errors.New("invalid slice permutation")
	}
	return res, nil
}

// SolvedPhase2Cube returns a solved Phase2Cube.
func SolvedPhase2Cube() Phase2Cube {
	return Phase2Cube{}
}

// Move applies a move to the Phase2Cube.
func (p *Phase2Cube) Move(move Phase2Move, table *Phase2Moves) {
	p.CornerPermutation = table.CornerMoves[p.CornerPermutation][int(move)]
	p.EdgePermutation = table.EdgeMoves[p.EdgePermutation][int(move)]
	p.SlicePermutation = table.SliceMoves[p.SlicePermutation][int(move)]
}

// Solved returns true if the Phase2Cube is solved.
func (p *Phase2Cube) Solved() bool {
	return p.CornerPermutation == 0 && p.EdgePermutation == 0 &&
		p.SlicePermutation == 0
}

// Phase2Move represents a move which can be applied to a Phase2Cube. This is a
// number in the range [0, 10), corresponding to F2 B2 R2 L2 U U' U2 D D' D2
// respectively.
type Phase2Move int

// Face returns a number from [1, 6] corresponding to the face of the move if it
// were applied on the Y axis.
func (p Phase2Move) Face() int {
	return []int{2, 3, 4, 5, 0, 0, 0, 1, 1, 1}[int(p)] + 1
}

// Inverse returns the move's inverse.
func (p Phase2Move) Inverse() Phase2Move {
	return []Phase2Move{0, 1, 2, 3, 5, 4, 6, 8, 7, 9}[int(p)]
}

// Move converts the Phase2Move into a regular Move.
// The axis argument indicates the axis that the move should act on (i.e. the
// axis of the corresponding Phase2Cube). This is a number in [0, 3).
func (p Phase2Move) Move(axis int) Move {
	return [][]Move{
		[]Move{14, 15, 12, 13, 5, 11, 17, 4, 10, 16},
		[]Move{14, 15, 16, 17, 0, 6, 12, 1, 7, 13},
		[]Move{13, 12, 16, 17, 2, 8, 14, 3, 9, 15},
	}[axis][int(p)]
}

// String returns the string representation of the move on the Y axis.
func (p Phase2Move) String() string {
	return p.Move(1).String()
}

// Phase2Moves is a table containing the necessary data to efficiently perform
// moves on a Phase2Cube.
type Phase2Moves struct {
	CornerMoves [40320][10]int
	EdgeMoves   [40320][10]int
	SliceMoves  [24][10]int
}

// NewPhase2Moves generates a Phase2Moves table.
func NewPhase2Moves() *Phase2Moves {
	res := new(Phase2Moves)

	perm8 := allPermutations(8)

	// We set some states to -1, that way we can tell which states have been
	// found and which have not.
	for i := 0; i < 40320; i++ {
		for j := 0; j < 10; j++ {
			res.CornerMoves[i][j] = -1
			res.EdgeMoves[i][j] = -1
		}
	}

	// Generate corner cases.
	for state, perm := range perm8 {
		corners := SolvedCubieCorners()
		// Permute the UD edges for the current case.
		for j, x := range perm {
			corners[j].Piece = x
		}
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			// The result of this move may have already been computed as an
			// inverse.
			if res.CornerMoves[state][m] >= 0 {
				continue
			}

			// Record the end state of this move.
			c := corners
			c.Move(Phase2Move(m).Move(1))
			endState := encodeYCornerPerm(&c)
			res.CornerMoves[state][m] = endState

			// For the end state, the inverse of this move gets the current
			// state.
			res.CornerMoves[endState][int(Phase2Move(m).Inverse())] = state
		}
	}

	// Generate edge cases.
	for state, perm := range perm8 {
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			// The result of this move may have already been computed as an
			// inverse.
			if res.EdgeMoves[state][m] >= 0 {
				continue
			}

			// Record the end state of this move.
			endState := moveUDEdgePerm(perm, Phase2Move(m))
			res.EdgeMoves[state][m] = endState

			// For the end state, the inverse of this move gets the current
			// state.
			res.EdgeMoves[endState][int(Phase2Move(m).Inverse())] = state
		}
	}

	// Generate slice moves
	for state, perm := range allPermutations(4) {
		edges := SolvedCubieEdges()
		// Permute the E slice for the current case.
		for j, x := range perm {
			slot := []int{1, 3, 7, 9}[j]
			piece := []int{1, 3, 7, 9}[x]
			edges[slot].Piece = piece
		}
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			e := edges
			e.Move(Phase2Move(m).Move(1))
			res.SliceMoves[state][m] = encodeESlicePerm(&e)
		}
	}

	return res
}

func encodeESlicePerm(e *CubieEdges) int {
	// Generate a permutation of {0, 1, 2, 3} that represents the permutation of
	// the E slice.
	var perm [4]int
	for i, slot := range []int{1, 3, 7, 9} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{-1, 0, -1, 1, -1, -1, -1, 2, -1, 3, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeFBEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{0, 1, 2, 3, 6, 7, 8, 9} {
		piece := (*e)[slot].Piece
		perm[i] = []int{0, 1, 2, 3, -1, -1, 4, 5, 6, 7, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeMSlicePerm(e *CubieEdges) int {
	// This is like encodeESlicePerm, but for the M slice.
	var perm [4]int
	for i, slot := range []int{0, 2, 6, 8} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{0, -1, 1, -1, -1, -1, 2, -1, 3, -1, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeRLEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{9, 4, 3, 10, 7, 5, 1, 11} {
		piece := (*e)[slot].Piece
		perm[i] = []int{-1, 6, -1, 2, 1, 5, -1, 4, -1, 0, 3, 7}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeSSlicePerm(e *CubieEdges) int {
	// This is like encodeESlicePerm, but for the S slice.
	var perm [4]int
	for i, slot := range []int{11, 10, 5, 4} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{-1, -1, -1, -1, 3, 2, -1, -1, -1, -1, 1, 0}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeUDEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{6, 5, 0, 4, 8, 11, 2, 10} {
		piece := (*e)[slot].Piece
		perm[i] = []int{2, -1, 6, -1, 3, 1, 0, -1, 4, -1, 7, 5}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeXCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i, idx := range xCornerIndices {
		perm[i] = inverseXCornerIndices[(*c)[idx].Piece]
	}
	return encodePermutationInPlace(perm[:])
}

func encodeYCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i := 0; i < 8; i++ {
		perm[i] = (*c)[i].Piece
	}
	return encodePermutationInPlace(perm[:])
}

func encodeZCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i, idx := range zCornerIndices {
		perm[i] = inverseZCornerIndices[(*c)[idx].Piece]
	}
	return encodePermutationInPlace(perm[:])
}

func moveYCornerPerm(perm []int, move Phase2Move) int {

}

func moveUDEdgePerm(perm []int, move Phase2Move) int {
	var newPerm [8]int
	copy(newPerm[:], perm)

	// The moves are F2, B2, R2, L2, U, U', U2, D, D', D2.
	switch (move) {
	case 0:
		newPerm[2], newPerm[6] = newPerm[6], newPerm[2]
	case 1:
		newPerm[0], newPerm[4] = newPerm[4], newPerm[0]
	case 2:
		newPerm[1], newPerm[5] = newPerm[5], newPerm[1]
	case 3:
		newPerm[3], newPerm[7] = newPerm[7], newPerm[3]
	case 4:
		newPerm[0], newPerm[1], newPerm[2], newPerm[3] =
			newPerm[3], newPerm[0], newPerm[1], newPerm[2]
	case 5:
		newPerm[3], newPerm[0], newPerm[1], newPerm[2] =
			newPerm[0], newPerm[1], newPerm[2], newPerm[3]
	case 6:
		newPerm[0], newPerm[2] = newPerm[2], newPerm[0]
		newPerm[1], newPerm[3] = newPerm[3], newPerm[1]
	case 7:
		newPerm[7], newPerm[4], newPerm[5], newPerm[6] =
			newPerm[4], newPerm[5], newPerm[6], newPerm[7]
	case 8:
		newPerm[4], newPerm[5], newPerm[6], newPerm[7] =
			newPerm[7], newPerm[4], newPerm[5], newPerm[6]
	case 9:
		newPerm[4], newPerm[6] = newPerm[6], newPerm[4]
		newPerm[5], newPerm[7] = newPerm[7], newPerm[5]
	}

	return encodePermutationInPlace(newPerm[:])
}
<commit_msg>faster Phase2Moves with moveYCornerPerm<commit_after>package gocube

import "errors"

// inverseXCornerIndices is the inverse permutation of xCornerIndices.
var inverseXCornerIndices []int = []int{2, 0, 3, 1, 6, 4, 7, 5}

// inverseZCornerIndices is the inverse permutation of zCornerIndices.
var inverseZCornerIndices []int = []int{4, 5, 0, 1, 6, 7, 2, 3}

// A Phase2Cube represents the parts of a cube that are important for phase-2
// solving.
type Phase2Cube struct {
	// CornerPermutation represents the permutation of the corners.
	CornerPermutation int

	// EdgePermutation represents the permutation of the 8 top/bottom edges.
	EdgePermutation int

	// SlicePermutation represents the permutation of the
	SlicePermutation int
}

// NewPhase2Cube generates a Phase2Cube from a CubieCube.
// The axis argument is 0 for X axis, 1 for Y axis, or 2 for Z axis.
// If the cube is not reduced to phase-2 in the given axis, this may return an
// error, but it does not validate everything.
func NewPhase2Cube(c CubieCube, axis int) (Phase2Cube, error) {
	var res Phase2Cube
	if axis == 0 {
		res.CornerPermutation = encodeXCornerPerm(&c.Corners)
		res.EdgePermutation = encodeRLEdges(&c.Edges)
		res.SlicePermutation = encodeMSlicePerm(&c.Edges)
	} else if axis == 1 {
		res.CornerPermutation = encodeYCornerPerm(&c.Corners)
		res.EdgePermutation = encodeUDEdges(&c.Edges)
		res.SlicePermutation = encodeESlicePerm(&c.Edges)
	} else {
		res.CornerPermutation = encodeZCornerPerm(&c.Corners)
		res.EdgePermutation = encodeFBEdges(&c.Edges)
		res.SlicePermutation = encodeSSlicePerm(&c.Edges)
	}
	if res.EdgePermutation < 0 {
		return res, errors.New("invalid edge permutation")
	} else if res.SlicePermutation < 0 {
		return res, errors.New("invalid slice permutation")
	}
	return res, nil
}

// SolvedPhase2Cube returns a solved Phase2Cube.
func SolvedPhase2Cube() Phase2Cube {
	return Phase2Cube{}
}

// Move applies a move to the Phase2Cube.
func (p *Phase2Cube) Move(move Phase2Move, table *Phase2Moves) {
	p.CornerPermutation = table.CornerMoves[p.CornerPermutation][int(move)]
	p.EdgePermutation = table.EdgeMoves[p.EdgePermutation][int(move)]
	p.SlicePermutation = table.SliceMoves[p.SlicePermutation][int(move)]
}

// Solved returns true if the Phase2Cube is solved.
func (p *Phase2Cube) Solved() bool {
	return p.CornerPermutation == 0 && p.EdgePermutation == 0 &&
		p.SlicePermutation == 0
}

// Phase2Move represents a move which can be applied to a Phase2Cube. This is a
// number in the range [0, 10), corresponding to F2 B2 R2 L2 U U' U2 D D' D2
// respectively.
type Phase2Move int

// Face returns a number from [1, 6] corresponding to the face of the move if it
// were applied on the Y axis.
func (p Phase2Move) Face() int {
	return []int{2, 3, 4, 5, 0, 0, 0, 1, 1, 1}[int(p)] + 1
}

// Inverse returns the move's inverse.
func (p Phase2Move) Inverse() Phase2Move {
	return []Phase2Move{0, 1, 2, 3, 5, 4, 6, 8, 7, 9}[int(p)]
}

// Move converts the Phase2Move into a regular Move.
// The axis argument indicates the axis that the move should act on (i.e. the
// axis of the corresponding Phase2Cube). This is a number in [0, 3).
func (p Phase2Move) Move(axis int) Move {
	return [][]Move{
		[]Move{14, 15, 12, 13, 5, 11, 17, 4, 10, 16},
		[]Move{14, 15, 16, 17, 0, 6, 12, 1, 7, 13},
		[]Move{13, 12, 16, 17, 2, 8, 14, 3, 9, 15},
	}[axis][int(p)]
}

// String returns the string representation of the move on the Y axis.
func (p Phase2Move) String() string {
	return p.Move(1).String()
}

// Phase2Moves is a table containing the necessary data to efficiently perform
// moves on a Phase2Cube.
type Phase2Moves struct {
	CornerMoves [40320][10]int
	EdgeMoves   [40320][10]int
	SliceMoves  [24][10]int
}

// NewPhase2Moves generates a Phase2Moves table.
func NewPhase2Moves() *Phase2Moves {
	res := new(Phase2Moves)

	perm8 := allPermutations(8)

	// We set some states to -1, that way we can tell which states have been
	// found and which have not.
	for i := 0; i < 40320; i++ {
		for j := 0; j < 10; j++ {
			res.CornerMoves[i][j] = -1
			res.EdgeMoves[i][j] = -1
		}
	}

	// Generate corner cases.
	for state, perm := range perm8 {
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			// The result of this move may have already been computed as an
			// inverse.
			if res.CornerMoves[state][m] >= 0 {
				continue
			}

			// Record the end state of this move.
			endState := moveYCornerPerm(perm, Phase2Move(m))
			res.CornerMoves[state][m] = endState

			// For the end state, the inverse of this move gets the current
			// state.
			res.CornerMoves[endState][int(Phase2Move(m).Inverse())] = state
		}
	}

	// Generate edge cases.
	for state, perm := range perm8 {
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			// The result of this move may have already been computed as an
			// inverse.
			if res.EdgeMoves[state][m] >= 0 {
				continue
			}

			// Record the end state of this move.
			endState := moveUDEdgePerm(perm, Phase2Move(m))
			res.EdgeMoves[state][m] = endState

			// For the end state, the inverse of this move gets the current
			// state.
			res.EdgeMoves[endState][int(Phase2Move(m).Inverse())] = state
		}
	}

	// Generate slice moves
	for state, perm := range allPermutations(4) {
		edges := SolvedCubieEdges()
		// Permute the E slice for the current case.
		for j, x := range perm {
			slot := []int{1, 3, 7, 9}[j]
			piece := []int{1, 3, 7, 9}[x]
			edges[slot].Piece = piece
		}
		// Apply all 10 moves to the cube.
		for m := 0; m < 10; m++ {
			e := edges
			e.Move(Phase2Move(m).Move(1))
			res.SliceMoves[state][m] = encodeESlicePerm(&e)
		}
	}

	return res
}

func encodeESlicePerm(e *CubieEdges) int {
	// Generate a permutation of {0, 1, 2, 3} that represents the permutation of
	// the E slice.
	var perm [4]int
	for i, slot := range []int{1, 3, 7, 9} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{-1, 0, -1, 1, -1, -1, -1, 2, -1, 3, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeFBEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{0, 1, 2, 3, 6, 7, 8, 9} {
		piece := (*e)[slot].Piece
		perm[i] = []int{0, 1, 2, 3, -1, -1, 4, 5, 6, 7, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeMSlicePerm(e *CubieEdges) int {
	// This is like encodeESlicePerm, but for the M slice.
	var perm [4]int
	for i, slot := range []int{0, 2, 6, 8} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{0, -1, 1, -1, -1, -1, 2, -1, 3, -1, -1, -1}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeRLEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{9, 4, 3, 10, 7, 5, 1, 11} {
		piece := (*e)[slot].Piece
		perm[i] = []int{-1, 6, -1, 2, 1, 5, -1, 4, -1, 0, 3, 7}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeSSlicePerm(e *CubieEdges) int {
	// This is like encodeESlicePerm, but for the S slice.
	var perm [4]int
	for i, slot := range []int{11, 10, 5, 4} {
		// Get the physical piece at the given slot.
		piece := (*e)[slot].Piece
		// Get an index from 0 to 3 for the piece.
		perm[i] = []int{-1, -1, -1, -1, 3, 2, -1, -1, -1, -1, 1, 0}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeUDEdges(e *CubieEdges) int {
	var perm [8]int
	for i, slot := range []int{6, 5, 0, 4, 8, 11, 2, 10} {
		piece := (*e)[slot].Piece
		perm[i] = []int{2, -1, 6, -1, 3, 1, 0, -1, 4, -1, 7, 5}[piece]
		if perm[i] < 0 {
			return -1
		}
	}
	return encodePermutationInPlace(perm[:])
}

func encodeXCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i, idx := range xCornerIndices {
		perm[i] = inverseXCornerIndices[(*c)[idx].Piece]
	}
	return encodePermutationInPlace(perm[:])
}

func encodeYCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i := 0; i < 8; i++ {
		perm[i] = (*c)[i].Piece
	}
	return encodePermutationInPlace(perm[:])
}

func encodeZCornerPerm(c *CubieCorners) int {
	var perm [8]int
	for i, idx := range zCornerIndices {
		perm[i] = inverseZCornerIndices[(*c)[idx].Piece]
	}
	return encodePermutationInPlace(perm[:])
}

func moveUDEdgePerm(perm []int, move Phase2Move) int {
	var newPerm [8]int
	copy(newPerm[:], perm)

	// The moves are F2, B2, R2, L2, U, U', U2, D, D', D2.
	switch (move) {
	case 0:
		newPerm[2], newPerm[6] = newPerm[6], newPerm[2]
	case 1:
		newPerm[0], newPerm[4] = newPerm[4], newPerm[0]
	case 2:
		newPerm[1], newPerm[5] = newPerm[5], newPerm[1]
	case 3:
		newPerm[3], newPerm[7] = newPerm[7], newPerm[3]
	case 4:
		newPerm[0], newPerm[1], newPerm[2], newPerm[3] =
			newPerm[3], newPerm[0], newPerm[1], newPerm[2]
	case 5:
		newPerm[3], newPerm[0], newPerm[1], newPerm[2] =
			newPerm[0], newPerm[1], newPerm[2], newPerm[3]
	case 6:
		newPerm[0], newPerm[2] = newPerm[2], newPerm[0]
		newPerm[1], newPerm[3] = newPerm[3], newPerm[1]
	case 7:
		newPerm[7], newPerm[4], newPerm[5], newPerm[6] =
			newPerm[4], newPerm[5], newPerm[6], newPerm[7]
	case 8:
		newPerm[4], newPerm[5], newPerm[6], newPerm[7] =
			newPerm[7], newPerm[4], newPerm[5], newPerm[6]
	case 9:
		newPerm[4], newPerm[6] = newPerm[6], newPerm[4]
		newPerm[5], newPerm[7] = newPerm[7], newPerm[5]
	}

	return encodePermutationInPlace(newPerm[:])
}

func moveYCornerPerm(perm []int, move Phase2Move) int {
	var p [8]int
	copy(p[:], perm)

	// The moves are F2, B2, R2, L2, U, U', U2, D, D', D2.
	switch (move) {
	case 0:
		p[5], p[6] = p[6], p[5]
		p[4], p[7] = p[7], p[4]
	case 1:
		p[1], p[2] = p[2], p[1]
		p[0], p[3] = p[3], p[0]
	case 2:
		p[1], p[7] = p[7], p[1]
		p[3], p[5] = p[5], p[3]
	case 3:
		p[0], p[6] = p[6], p[0]
		p[2], p[4] = p[4], p[2]
	case 4:
		p[2], p[3], p[7], p[6] = p[6], p[2], p[3], p[7]
	case 5:
		p[6], p[2], p[3], p[7] = p[2], p[3], p[7], p[6]
	case 6:
		p[2], p[7] = p[7], p[2]
		p[3], p[6] = p[6], p[3]
	case 7:
		p[0], p[1], p[5], p[4] = p[1], p[5], p[4], p[0]
	case 8:
		p[1], p[5], p[4], p[0] = p[0], p[1], p[5], p[4]
	case 9:
		p[0], p[5] = p[5], p[0]
		p[1], p[4] = p[4], p[1]
	}

	return encodePermutationInPlace(p[:])
}
<|endoftext|>


<commit_before>package pig

import (
	"crypto/sha1"
	"fmt"
	"io"
	"strconv"
	"strings"
)

// StringSliceContains determines whether this slice contain this string?
func StringSliceContains(hay []string, needle string) bool {
	for _, straw := range hay {
		if straw == needle {
			return true
		}
	}
	return false
}

// StringSliceEquals determines whether two string slices are the same.
// Arbitrary regards nil and []string{} as NOT equal.
func StringSliceEquals(a []string, b []string) bool {
	if (a == nil) != (b == nil) {
		return false
	}
	n := len(a)
	if n != len(b) {
		return false
	}
	for i := 0; i < n; i++ {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// Last returns the last element of a slice,
// or empty string if the slice is empty.
func Last(array []string) (result string) {
	if n := len(array); n >= 1 {
		result = array[n-1]
	}
	return
}

// FilterOut creates a new slice from s, removing any element contained in forbidden.
// Accept nils and empty lists.
// Never yields nil, but may yield an empty list.
func FilterOut(s []string, forbidden []string) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if !StringSliceContains(forbidden, v) {
			p = append(p, v)
		}
	}
	return p
}

// FilterStrings creates a new slice from s, removing any element that doesn't match the predicate fn.
//
// From http://blog.golang.org/go-slices-usage-and-internals
//
// Accept nils and empty lists.
// Never yields nil, but may yield an empty list.
//
// Items that match the predicate are kept (not thrown away, do not confuse!)
func FilterStrings(s []string, fn func(string) bool) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if fn(v) {
			p = append(p, v)
		}
	}
	return p
}

// RemoveEmptyStrings creates a new slice from s, removing all empty elements.
func RemoveEmptyStrings(s []string) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if v != "" {
			p = append(p, v)
		}
	}
	return p
}

// MapStrings applies f to each element of str, and returns a new slice of the
// same size containing all the results.
func MapStrings(str []string, f func(string) string) []string {
	result := make([]string, len(str))
	for i, s := range str {
		result[i] = f(s)
	}
	return result
}

// Sha1hash is a shorthand to get a sha1 hash as a string.
func Sha1hash(s string) string {
	h := sha1.New()
	io.WriteString(h, s)
	return fmt.Sprintf("%x", h.Sum(nil))
}

// String2Int is a shorthand to convert a string to an int.
// It doesn't panic and does'nt return any error, but it may return -1.
// Thus it should be used only to parse positive integer that are extermely
// likely to be well-formed.
func String2Int(str string) int {
	if number64, err := strconv.ParseInt(str, 10, 32); err == nil {
		return int(number64)
	}
	return -1
}

// Min returns the minimum of a and b.
func Min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Max returns the maximum of a and b.
func Max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// Shorten keeps only the n first bytes, add appends "..." if needed.
func Shorten(s string, n int) string {
	// TODO better handle UTF-8
	if len(s) > n {
		return s[:n] + "..."
	}
	return s
}

// Flatten replaces newlines by spaces
func Flatten(s string) string {
	s = strings.Replace(s, "\r\n", " ", -1)
	s = strings.Replace(s, "\n\r", " ", -1)
	s = strings.Replace(s, "\n", " ", -1)
	return s
}

func Truncate(s string, maxChars int) string {
	runes := []rune(s)
	if len(runes) < maxChars {
		return s
	}
	return string(runes[:maxChars])
}

// Concurrent launches provided funcs, and waits for their completion.
func Concurrent(funcs ...func()) {
	var wg sync.WaitGroup
	wg.Add(len(funcs))
	for _, f := range funcs {
		f := f
		go func() {
			f()
			wg.Done()
		}()
	}
	wg.Wait()
}

// Concurrent launches provided funcs, and returns a channel to notify completion.
func ConcurrentPromise(funcs ...func()) chan bool {
	ch := make(chan bool)
	go func() {
		Concurrent(funcs...)
		ch <- true
		close(ch)
	}()
	return ch
}
<commit_msg>Util func ConcurrentWithErrors.<commit_after>package pig

import (
	"crypto/sha1"
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
)

// StringSliceContains determines whether this slice contain this string?
func StringSliceContains(hay []string, needle string) bool {
	for _, straw := range hay {
		if straw == needle {
			return true
		}
	}
	return false
}

// StringSliceEquals determines whether two string slices are the same.
// Arbitrary regards nil and []string{} as NOT equal.
func StringSliceEquals(a []string, b []string) bool {
	if (a == nil) != (b == nil) {
		return false
	}
	n := len(a)
	if n != len(b) {
		return false
	}
	for i := 0; i < n; i++ {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// Last returns the last element of a slice,
// or empty string if the slice is empty.
func Last(array []string) (result string) {
	if n := len(array); n >= 1 {
		result = array[n-1]
	}
	return
}

// FilterOut creates a new slice from s, removing any element contained in forbidden.
// Accept nils and empty lists.
// Never yields nil, but may yield an empty list.
func FilterOut(s []string, forbidden []string) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if !StringSliceContains(forbidden, v) {
			p = append(p, v)
		}
	}
	return p
}

// FilterStrings creates a new slice from s, removing any element that doesn't match the predicate fn.
//
// From http://blog.golang.org/go-slices-usage-and-internals
//
// Accept nils and empty lists.
// Never yields nil, but may yield an empty list.
//
// Items that match the predicate are kept (not thrown away, do not confuse!)
func FilterStrings(s []string, fn func(string) bool) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if fn(v) {
			p = append(p, v)
		}
	}
	return p
}

// RemoveEmptyStrings creates a new slice from s, removing all empty elements.
func RemoveEmptyStrings(s []string) []string {
	p := make([]string, 0, len(s))
	for _, v := range s {
		if v != "" {
			p = append(p, v)
		}
	}
	return p
}

// MapStrings applies f to each element of str, and returns a new slice of the
// same size containing all the results.
func MapStrings(str []string, f func(string) string) []string {
	result := make([]string, len(str))
	for i, s := range str {
		result[i] = f(s)
	}
	return result
}

// Sha1hash is a shorthand to get a sha1 hash as a string.
func Sha1hash(s string) string {
	h := sha1.New()
	io.WriteString(h, s)
	return fmt.Sprintf("%x", h.Sum(nil))
}

// String2Int is a shorthand to convert a string to an int.
// It doesn't panic and does'nt return any error, but it may return -1.
// Thus it should be used only to parse positive integer that are extermely
// likely to be well-formed.
func String2Int(str string) int {
	if number64, err := strconv.ParseInt(str, 10, 32); err == nil {
		return int(number64)
	}
	return -1
}

// Min returns the minimum of a and b.
func Min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Max returns the maximum of a and b.
func Max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// Shorten keeps only the n first bytes, add appends "..." if needed.
func Shorten(s string, n int) string {
	// TODO better handle UTF-8
	if len(s) > n {
		return s[:n] + "..."
	}
	return s
}

// Flatten replaces newlines by spaces
func Flatten(s string) string {
	s = strings.Replace(s, "\r\n", " ", -1)
	s = strings.Replace(s, "\n\r", " ", -1)
	s = strings.Replace(s, "\n", " ", -1)
	return s
}

func Truncate(s string, maxChars int) string {
	runes := []rune(s)
	if len(runes) < maxChars {
		return s
	}
	return string(runes[:maxChars])
}

// Concurrent launches provided funcs, and waits for their completion.
func Concurrent(funcs ...func()) {
	var wg sync.WaitGroup
	wg.Add(len(funcs))
	for _, f := range funcs {
		f := f
		go func() {
			f()
			wg.Done()
		}()
	}
	wg.Wait()
}

// Concurrent launches provided funcs, and returns a channel to notify completion.
func ConcurrentPromise(funcs ...func()) chan bool {
	ch := make(chan bool)
	go func() {
		Concurrent(funcs...)
		ch <- true
		close(ch)
	}()
	return ch
}

// ConcurrentWithErrors launches provided funcs, and gathers errors.
// If no errors, ok is true and the returned slice contains all nil values.
func ConcurrentWithErrors(funcs ...func() error) (ok bool, errs []error) {
	errs = make([]error, len(funcs))
	var wg sync.WaitGroup
	wg.Add(len(funcs))
	for i, f := range funcs {
		i := i
		f := f
		go func() {
			errs[i] = f()
			wg.Done()
		}()
	}
	wg.Wait()
	ok = true
	for _, err := range errs {
		if err != nil {
			ok = false
		}
	}
	return ok, errs
}
<|endoftext|>


<commit_before>package cli

import (
	"github.com/containerum/chkit/pkg/cli/deployment"
	"github.com/containerum/chkit/pkg/cli/prerun"
	"github.com/containerum/chkit/pkg/cli/solution"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/util/angel"
	"github.com/containerum/chkit/pkg/util/ferr"
	"github.com/spf13/cobra"
)

func Run(ctx *context.Context) *cobra.Command {
	command := &cobra.Command{
		Use:   "run",
		Short: "Run solutions and pods",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if err := prerun.PreRun(ctx); err != nil {
				angel.Angel(ctx, err)
				ctx.Exit(1)
			}
			if err := prerun.GetNamespaceByUserfriendlyID(ctx, cmd.Flags()); err != nil {
				ferr.Println(err)
				ctx.Exit(1)
			}
		},
		Run: func(command *cobra.Command, args []string) {
			command.Help()
		},
		PersistentPostRun: ctx.CobraPostRun,
	}
	command.AddCommand(
		clisolution.Run(ctx),
		clideployment.RunVersion(ctx),
	)
	command.PersistentFlags().
		StringP("namespace", "n", ctx.GetNamespace().ID, "")
	command.PersistentFlags().
		BoolP("help", "h", false, "Print help for chkit")
	return command
}
<commit_msg>fix short help message<commit_after>package cli

import (
	"github.com/containerum/chkit/pkg/cli/deployment"
	"github.com/containerum/chkit/pkg/cli/prerun"
	"github.com/containerum/chkit/pkg/cli/solution"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/util/angel"
	"github.com/containerum/chkit/pkg/util/ferr"
	"github.com/spf13/cobra"
)

func Run(ctx *context.Context) *cobra.Command {
	command := &cobra.Command{
		Use:   "run",
		Short: "Run solutions and deployments",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if err := prerun.PreRun(ctx); err != nil {
				angel.Angel(ctx, err)
				ctx.Exit(1)
			}
			if err := prerun.GetNamespaceByUserfriendlyID(ctx, cmd.Flags()); err != nil {
				ferr.Println(err)
				ctx.Exit(1)
			}
		},
		Run: func(command *cobra.Command, args []string) {
			command.Help()
		},
		PersistentPostRun: ctx.CobraPostRun,
	}
	command.AddCommand(
		clisolution.Run(ctx),
		clideployment.RunVersion(ctx),
	)
	command.PersistentFlags().
		StringP("namespace", "n", ctx.GetNamespace().ID, "")
	command.PersistentFlags().
		BoolP("help", "h", false, "Print help for chkit")
	return command
}
<|endoftext|>


<commit_before>package pkg

import (
	"errors"
	"strings"
	"sync"

	log "github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

//IngressList ...
type IngressList struct {
	Name           string
	Namespace      string
	Context        string
	Host           string
	Class          string
	WhitelistedIps []string
}

//MyClientSet ...
type MyClientSet struct {
	Clientset *kubernetes.Clientset
	sync.RWMutex
}

//GetIngresses ...
func (c MyClientSet) GetIngresses(context string, ns string) ([]IngressList, error) {
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingressLists, err := ingressClient.List(metav1.ListOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return nil, err
	}

	var myIngress []IngressList

	for _, ingress := range ingressLists.Items {
		if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {
			var ingressHosts string
			if _, ok := ingress.Annotations["concierge-url"]; ok {
				ingressHosts = ingress.Annotations["concierge-url"]
			} else {
				for _, hosts := range ingress.Spec.Rules {
					if ingressHosts != "" {
						ingressHosts = ingressHosts + ", " + hosts.Host
					} else {
						ingressHosts = hosts.Host
					}
				}
			}

			myIngress = append(myIngress, IngressList{
				ingress.Name,
				ingress.Namespace,
				context,
				ingressHosts,
				ingress.Annotations["kubernetes.io/ingress.class"],
				strings.Split(ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"], ","),
			})
		}
	}
	return myIngress, nil
}

//RemoveIngressIP ...
func (c MyClientSet) RemoveIngressIP(ns string, ingressName string, ip string) (bool, bool, error) {
	c.Lock()
	defer c.Unlock()
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return false, false, err
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {

		whitelistIps, ok := ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"]
		if ok {
			whitelistIpsArr := strings.Split(whitelistIps, ",")
			for i := range whitelistIpsArr {
				whitelistIpsArr[i] = strings.TrimSpace(whitelistIpsArr[i])
			}
			var newWhitelistIpsArr []string
			for _, ips := range whitelistIpsArr {
				if ips != ip {
					newWhitelistIpsArr = append(newWhitelistIpsArr, ips)
				}
			}
			whitelistIps = strings.Join(newWhitelistIpsArr, ", ")
			ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"] = whitelistIps
			_, updateErr := ingressClient.Update(ingress)

			if updateErr != nil {
				log.Error("Error: ", updateErr)
				return false, true, updateErr
			}
			return true, false, nil
		}
	}
	err = errors.New("Ingress Annotation not found")
	log.Error("Error: ", err)
	return false, false, err
}

//WhiteListIP ...
func (c MyClientSet) WhiteListIP(ns string, ingressName string, ip string) (bool, error) {
	c.Lock()
	defer c.Unlock()
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return false, err
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {

		whitelistIps, ok := ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"]
		if ok {
			whitelistIpsArr := strings.Split(whitelistIps, ",")
			for i := range whitelistIpsArr {
				whitelistIpsArr[i] = strings.TrimSpace(whitelistIpsArr[i])
			}
			new := true
			for _, ips := range whitelistIpsArr {
				if ips == ip {
					log.Warn("Your IP is already present there")
					new = false
					break
				}
			}
			if new {
				whitelistIpsArr = append(whitelistIpsArr, ip)
				whitelistIps = strings.Join(whitelistIpsArr, ", ")
				ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"] = whitelistIps
				_, updateErr := ingressClient.Update(ingress)

				if updateErr != nil {
					log.Error("Error: ", updateErr)
					return false, updateErr
				}
				return true, nil
			}
			return false, nil
		}
	}
	err = errors.New("Ingress Annotation not found")
	log.Error("Error: ", err)
	return false, err
}

//GetIngress ...
func (c MyClientSet) GetIngress(context string, ns string, ingressName string) (IngressList, error) {
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)
	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return IngressList{}, nil
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {
		var ingressHosts string
		if _, ok := ingress.Annotations["concierge-url"]; ok {
			ingressHosts = ingress.Annotations["concierge-url"]
		} else {
			for _, hosts := range ingress.Spec.Rules {
				if ingressHosts != "" {
					ingressHosts = ingressHosts + ", " + hosts.Host
				} else {
					ingressHosts = hosts.Host
				}
			}
		}
		myIngress := IngressList{
			ingress.Name,
			ingress.Namespace,
			context,
			ingressHosts,
			ingress.Annotations["kubernetes.io/ingress.class"],
			[]string{},
		}
		return myIngress, nil
	}
	log.Error("Error: ", "Ingress Annotation not found")
	return IngressList{}, nil
}
<commit_msg>Fixed minor err issue<commit_after>package pkg

import (
	"errors"
	"strings"
	"sync"

	log "github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

//IngressList ...
type IngressList struct {
	Name           string
	Namespace      string
	Context        string
	Host           string
	Class          string
	WhitelistedIps []string
}

//MyClientSet ...
type MyClientSet struct {
	Clientset *kubernetes.Clientset
	sync.RWMutex
}

//GetIngresses ...
func (c MyClientSet) GetIngresses(context string, ns string) ([]IngressList, error) {
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingressLists, err := ingressClient.List(metav1.ListOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return nil, err
	}

	var myIngress []IngressList

	for _, ingress := range ingressLists.Items {
		if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {
			var ingressHosts string
			if _, ok := ingress.Annotations["concierge-url"]; ok {
				ingressHosts = ingress.Annotations["concierge-url"]
			} else {
				for _, hosts := range ingress.Spec.Rules {
					if ingressHosts != "" {
						ingressHosts = ingressHosts + ", " + hosts.Host
					} else {
						ingressHosts = hosts.Host
					}
				}
			}

			myIngress = append(myIngress, IngressList{
				ingress.Name,
				ingress.Namespace,
				context,
				ingressHosts,
				ingress.Annotations["kubernetes.io/ingress.class"],
				strings.Split(ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"], ","),
			})
		}
	}
	return myIngress, nil
}

//RemoveIngressIP ...
func (c MyClientSet) RemoveIngressIP(ns string, ingressName string, ip string) (bool, bool, error) {
	c.Lock()
	defer c.Unlock()
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return false, false, err
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {

		whitelistIps, ok := ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"]
		if ok {
			whitelistIpsArr := strings.Split(whitelistIps, ",")
			for i := range whitelistIpsArr {
				whitelistIpsArr[i] = strings.TrimSpace(whitelistIpsArr[i])
			}
			var newWhitelistIpsArr []string
			for _, ips := range whitelistIpsArr {
				if ips != ip {
					newWhitelistIpsArr = append(newWhitelistIpsArr, ips)
				}
			}
			whitelistIps = strings.Join(newWhitelistIpsArr, ", ")
			ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"] = whitelistIps
			_, updateErr := ingressClient.Update(ingress)

			if updateErr != nil {
				log.Error("Error: ", updateErr)
				return false, true, updateErr
			}
			return true, false, nil
		}
	}
	err = errors.New("Ingress Annotation not found")
	log.Error("Error: ", err)
	return false, false, err
}

//WhiteListIP ...
func (c MyClientSet) WhiteListIP(ns string, ingressName string, ip string) (bool, error) {
	c.Lock()
	defer c.Unlock()
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)

	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return false, err
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {

		whitelistIps, ok := ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"]
		if ok {
			whitelistIpsArr := strings.Split(whitelistIps, ",")
			for i := range whitelistIpsArr {
				whitelistIpsArr[i] = strings.TrimSpace(whitelistIpsArr[i])
			}
			new := true
			for _, ips := range whitelistIpsArr {
				if ips == ip {
					log.Warn("Your IP is already present there")
					new = false
					break
				}
			}
			if new {
				whitelistIpsArr = append(whitelistIpsArr, ip)
				whitelistIps = strings.Join(whitelistIpsArr, ", ")
				ingress.Annotations["traefik.ingress.kubernetes.io/whitelist-source-range"] = whitelistIps
				_, updateErr := ingressClient.Update(ingress)

				if updateErr != nil {
					log.Error("Error: ", updateErr)
					return false, updateErr
				}
				return true, nil
			}
			return false, nil
		}
	}
	err = errors.New("Ingress Annotation not found")
	log.Error("Error: ", err)
	return false, err
}

//GetIngress ...
func (c MyClientSet) GetIngress(context string, ns string, ingressName string) (IngressList, error) {
	ingressClient := c.Clientset.ExtensionsV1beta1().Ingresses(ns)
	ingress, err := ingressClient.Get(ingressName, metav1.GetOptions{})
	if err != nil {
		log.Error("Error: ", err)
		return IngressList{}, err
	}
	if _, ok := ingress.Annotations["concierge"]; ok && ingress.Annotations["concierge"] == "true" {
		var ingressHosts string
		if _, ok := ingress.Annotations["concierge-url"]; ok {
			ingressHosts = ingress.Annotations["concierge-url"]
		} else {
			for _, hosts := range ingress.Spec.Rules {
				if ingressHosts != "" {
					ingressHosts = ingressHosts + ", " + hosts.Host
				} else {
					ingressHosts = hosts.Host
				}
			}
		}
		myIngress := IngressList{
			ingress.Name,
			ingress.Namespace,
			context,
			ingressHosts,
			ingress.Annotations["kubernetes.io/ingress.class"],
			[]string{},
		}
		return myIngress, nil
	}
	err = errors.New("Ingress Annotation not found")
	log.Error("Error: ", err)
	return IngressList{}, err
}
<|endoftext|>