<commit_before>package awsprovider

import (
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/stackplan"
	"koding/kites/kloud/terraformer"
	tf "koding/kites/terraformer"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	awssession "github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"golang.org/x/net/context"
	"gopkg.in/mgo.v2/bson"
)

// Bootstrap
func (s *Stack) Bootstrap(ctx context.Context) (interface{}, error) {
	var arg kloud.BootstrapRequest
	if err := s.Req.Args.One().Unmarshal(&arg); err != nil {
		return nil, err
	}

	if err := arg.Valid(); err != nil {
		return nil, err
	}

	if arg.Destroy {
		s.Log.Debug("Bootstrap destroy is called")
	} else {
		s.Log.Debug("Bootstrap apply is called")
	}

	if err := s.Builder.BuildCredentials(s.Req.Method, s.Req.Username, arg.GroupName, arg.Identifiers); err != nil {
		return nil, err
	}

	s.Log.Debug("Connecting to terraformer kite")

	tfKite, err := terraformer.Connect(s.Session.Kite)
	if err != nil {
		return nil, err
	}
	defer tfKite.Close()

	s.Log.Debug("Iterating over credentials")

	for _, cred := range s.Builder.Credentials {
		// We are going to support more providers in the future, for now only allow aws
		if cred.Provider != "aws" {
			return nil, fmt.Errorf("Bootstrap is only supported for 'aws' provider. Got: '%s'", cred.Provider)
		}

		meta := cred.Meta.(*AwsMeta)

		sess := awssession.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(meta.AccessKey, meta.SecretKey, ""),
			Region:      aws.String(meta.Region),
		})

		iamClient := iam.New(sess)

		s.Log.Debug("Fetching the AWS user information to get the account ID")
		user, err := iamClient.GetUser(nil) // will default to username making the request
		if err != nil {
			return nil, err
		}

		awsAccountID, err := stackplan.ParseAccountID(aws.StringValue(user.User.Arn))
		if err != nil {
			return nil, err
		}

		contentID := fmt.Sprintf("%s-%s-%s", awsAccountID, arg.GroupName, cred.Identifier)
		s.Log.Debug("Building template: %s", contentID)

		if err := s.Builder.BuildTemplate(awsBootstrap, contentID); err != nil {
			return nil, err
		}

		keyName := "koding-deployment-" + s.Req.Username + "-" + strconv.FormatInt(time.Now().UTC().UnixNano(), 10)

		finalBootstrap, err := s.Builder.Template.JsonOutput()
		if err != nil {
			return nil, err
		}

		finalBootstrap, err = appendAWSTemplateData(finalBootstrap, &awsTemplateData{
			KeyPairName:     keyName,
			PublicKey:       s.Keys.PublicKey,
			EnvironmentName: fmt.Sprintf("Koding-%s-Bootstrap", arg.GroupName),
		})

		s.Log.Debug("Final bootstrap template:")
		s.Log.Debug(finalBootstrap)

		if err != nil {
			return nil, err
		}

		// Important so bootstraping is distributed amongs multiple users. If I
		// use these keys to bootstrap, any other user should be not create
		// again, instead they should be fetch and use the existing bootstrap
		// data.

		resp := &AwsMeta{}

		// this is custom because we need to remove the fields if we get a
		// destroy. So the operator changes from $set to $unset.
		mongodDBOperator := "$set"

		if arg.Destroy {
			mongodDBOperator = "$unset"
			s.Log.Info("Destroying bootstrap resources belonging to identifier '%s'", cred.Identifier)
			_, err := tfKite.Destroy(&tf.TerraformRequest{
				Content:   finalBootstrap,
				ContentID: contentID,
				TraceID:   s.TraceID,
			})
			if err != nil {
				return nil, err
			}
		} else {
			s.Log.Info("Creating bootstrap resources belonging to identifier '%s'", cred.Identifier)
			state, err := tfKite.Apply(&tf.TerraformRequest{
				Content:   finalBootstrap,
				ContentID: contentID,
				TraceID:   s.TraceID,
			})
			if err != nil {
				return nil, err
			}

			s.Log.Debug("[%s] state.RootModule().Outputs = %+v\n", cred.Identifier, state.RootModule().Outputs)

			if err := s.Builder.Object.Decode(state.RootModule().Outputs, resp); err != nil {
				return nil, err
			}

			s.Log.Debug("[%s] resp = %+v\n", cred.Identifier, resp)

			if err := resp.BootstrapValid(); err != nil {
				return nil, fmt.Errorf("invalid bootstrap metadata for %q: %s", cred.Identifier, err)
			}
		}

		s.Log.Debug("[%s] Bootstrap response: %+v", cred.Identifier, resp)

		if err := modelhelper.UpdateCredentialData(cred.Identifier, bson.M{
			mongodDBOperator: bson.M{
				"meta.acl":        resp.ACL,
				"meta.cidr_block": resp.CidrBlock,
				"meta.igw":        resp.IGW,
				"meta.key_pair":   resp.KeyPair,
				"meta.rtb":        resp.RTB,
				"meta.sg":         resp.SG,
				"meta.subnet":     resp.Subnet,
				"meta.vpc":        resp.VPC,
				"meta.ami":        resp.AMI,
			},
		}); err != nil {
			return nil, err
		}
	}

	return true, nil
}

func appendAWSTemplateData(template string, awsData *awsTemplateData) (string, error) {
	var data struct {
		Output   map[string]map[string]interface{} `json:"output,omitempty"`
		Resource map[string]map[string]interface{} `json:"resource,omitempty"`
		Provider map[string]map[string]interface{} `json:"provider,omitempty"`
		Variable map[string]map[string]interface{} `json:"variable,omitempty"`
	}

	if err := json.Unmarshal([]byte(template), &data); err != nil {
		return "", err
	}

	data.Variable["key_name"] = map[string]interface{}{
		"default": awsData.KeyPairName,
	}

	data.Variable["public_key"] = map[string]interface{}{
		"default": awsData.PublicKey,
	}

	data.Variable["environment_name"] = map[string]interface{}{
		"default": awsData.EnvironmentName,
	}

	out, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return "", err
	}

	return string(out), nil
}

// awsTemplateData is being used to format the bootstrap before we pass it to
// terraformer
type awsTemplateData struct {
	KeyPairName     string
	PublicKey       string
	EnvironmentName string
}

var awsBootstrap = `{
    "provider": {
        "aws": {
            "access_key": "${var.aws_access_key}",
            "secret_key": "${var.aws_secret_key}",
            "region": "${var.aws_region}"
        }
    },
    "output": {
        "vpc": {
            "value": "${aws_vpc.vpc.id}"
        },
        "cidr_block": {
            "value": "${aws_vpc.vpc.cidr_block}"
        },
        "rtb": {
            "value": "${aws_vpc.vpc.main_route_table_id}"
        },
        "acl": {
            "value": "${aws_vpc.vpc.default_network_acl_id}"
        },
        "igw": {
            "value": "${aws_internet_gateway.main_vpc_igw.id}"
        },
        "subnet": {
            "value": "${aws_subnet.main_koding_subnet.id}"
        },
        "sg": {
            "value": "${aws_security_group.allow_all.id}"
        },
        "ami": {
            "value": "${lookup(var.aws_amis, var.aws_region)}"
        },
        "key_pair": {
            "value": "${aws_key_pair.koding_key_pair.key_name}"
        }
    },
    "resource": {
        "aws_vpc": {
            "vpc": {
                "cidr_block": "${var.cidr_block}",
                "tags": {
                    "Name": "${var.environment_name}"
                }
            }
        },
        "aws_internet_gateway": {
            "main_vpc_igw": {
                "tags": {
                    "Name": "${var.environment_name}"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_subnet": {
            "main_koding_subnet": {
                "availability_zone": "${lookup(var.aws_availability_zones, var.aws_region)}",
                "cidr_block": "${var.cidr_block}",
                "map_public_ip_on_launch": true,
                "tags": {
                    "Name": "${var.environment_name}",
                    "subnet": "public"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_route_table": {
            "public": {
                "route": {
                    "cidr_block": "0.0.0.0/0",
                    "gateway_id": "${aws_internet_gateway.main_vpc_igw.id}"
                },
                "tags": {
                    "Name": "${var.environment_name}",
                    "routeTable": "koding",
                    "subnet": "public"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_route_table_association": {
            "public-1": {
                "route_table_id": "${aws_route_table.public.id}",
                "subnet_id": "${aws_subnet.main_koding_subnet.id}"
            }
        },
        "aws_security_group": {
            "allow_all": {
                "description": "Allow all inbound and outbound traffic",
                "ingress": {
                    "from_port": 0,
                    "to_port": 0,
                    "protocol": "-1",
                    "cidr_blocks": ["0.0.0.0/0"],
                    "self": true
                },
                "egress": {
                    "from_port": 0,
                    "to_port": 0,
                    "protocol": "-1",
                    "self": true,
                    "cidr_blocks": ["0.0.0.0/0"]
                },
                "name": "allow_all",
                "tags": {
                    "Name": "${var.environment_name}"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_key_pair": {
            "koding_key_pair": {
                "key_name": "${var.key_name}",
                "public_key": "${var.public_key}"
            }
        }
    },
    "variable": {
        "cidr_block": {
            "default": "10.0.0.0/16"
        },
        "environment_name": {
            "default": "Koding-Bootstrap"
        },
        "aws_availability_zones": {
            "default": {
                "ap-northeast-1": "ap-northeast-1b",
                "ap-southeast-1": "ap-southeast-1b",
                "ap-southeast-2": "ap-southeast-2b",
                "eu-central-1": "eu-central-1b",
                "eu-west-1": "eu-west-1b",
                "sa-east-1": "sa-east-1b",
                "us-east-1": "us-east-1b",
                "us-west-1": "us-west-1b",
                "us-west-2": "us-west-2b"
            }
        },
        "aws_amis": {
            "default": {
                "ap-northeast-1": "ami-9e5cff9e",
                "ap-southeast-1": "ami-ec7879be",
                "ap-southeast-2": "ami-2fce8b15",
                "eu-central-1": "ami-60f9c27d",
                "eu-west-1": "ami-7c4b0a0b",
                "sa-east-1": "ami-cd9518d0",
                "us-east-1": "ami-cf35f3a4",
                "us-west-1": "ami-b33dccf7",
                "us-west-2": "ami-8d5b5dbd"
            }
        }
    }
}`
<commit_msg>Bootstrap: include group slug in key pair names while bootstrapping<commit_after>package awsprovider

import (
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/stackplan"
	"koding/kites/kloud/terraformer"
	tf "koding/kites/terraformer"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	awssession "github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"golang.org/x/net/context"
	"gopkg.in/mgo.v2/bson"
)

// Bootstrap
func (s *Stack) Bootstrap(ctx context.Context) (interface{}, error) {
	var arg kloud.BootstrapRequest
	if err := s.Req.Args.One().Unmarshal(&arg); err != nil {
		return nil, err
	}

	if err := arg.Valid(); err != nil {
		return nil, err
	}

	if arg.Destroy {
		s.Log.Debug("Bootstrap destroy is called")
	} else {
		s.Log.Debug("Bootstrap apply is called")
	}

	if err := s.Builder.BuildCredentials(s.Req.Method, s.Req.Username, arg.GroupName, arg.Identifiers); err != nil {
		return nil, err
	}

	s.Log.Debug("Connecting to terraformer kite")

	tfKite, err := terraformer.Connect(s.Session.Kite)
	if err != nil {
		return nil, err
	}
	defer tfKite.Close()

	s.Log.Debug("Iterating over credentials")

	for _, cred := range s.Builder.Credentials {
		// We are going to support more providers in the future, for now only allow aws
		if cred.Provider != "aws" {
			return nil, fmt.Errorf("Bootstrap is only supported for 'aws' provider. Got: '%s'", cred.Provider)
		}

		meta := cred.Meta.(*AwsMeta)

		sess := awssession.New(&aws.Config{
			Credentials: credentials.NewStaticCredentials(meta.AccessKey, meta.SecretKey, ""),
			Region:      aws.String(meta.Region),
		})

		iamClient := iam.New(sess)

		s.Log.Debug("Fetching the AWS user information to get the account ID")
		user, err := iamClient.GetUser(nil) // will default to username making the request
		if err != nil {
			return nil, err
		}

		awsAccountID, err := stackplan.ParseAccountID(aws.StringValue(user.User.Arn))
		if err != nil {
			return nil, err
		}

		contentID := fmt.Sprintf("%s-%s-%s", awsAccountID, arg.GroupName, cred.Identifier)
		s.Log.Debug("Building template: %s", contentID)

		if err := s.Builder.BuildTemplate(awsBootstrap, contentID); err != nil {
			return nil, err
		}

		keyName := "koding-deployment-" + s.Req.Username + "-" + arg.GroupName + "-" + strconv.FormatInt(time.Now().UTC().UnixNano(), 10)

		finalBootstrap, err := s.Builder.Template.JsonOutput()
		if err != nil {
			return nil, err
		}

		finalBootstrap, err = appendAWSTemplateData(finalBootstrap, &awsTemplateData{
			KeyPairName:     keyName,
			PublicKey:       s.Keys.PublicKey,
			EnvironmentName: fmt.Sprintf("Koding-%s-Bootstrap", arg.GroupName),
		})

		s.Log.Debug("Final bootstrap template:")
		s.Log.Debug(finalBootstrap)

		if err != nil {
			return nil, err
		}

		// Important so bootstraping is distributed amongs multiple users. If I
		// use these keys to bootstrap, any other user should be not create
		// again, instead they should be fetch and use the existing bootstrap
		// data.

		resp := &AwsMeta{}

		// this is custom because we need to remove the fields if we get a
		// destroy. So the operator changes from $set to $unset.
		mongodDBOperator := "$set"

		if arg.Destroy {
			mongodDBOperator = "$unset"
			s.Log.Info("Destroying bootstrap resources belonging to identifier '%s'", cred.Identifier)
			_, err := tfKite.Destroy(&tf.TerraformRequest{
				Content:   finalBootstrap,
				ContentID: contentID,
				TraceID:   s.TraceID,
			})
			if err != nil {
				return nil, err
			}
		} else {
			s.Log.Info("Creating bootstrap resources belonging to identifier '%s'", cred.Identifier)
			state, err := tfKite.Apply(&tf.TerraformRequest{
				Content:   finalBootstrap,
				ContentID: contentID,
				TraceID:   s.TraceID,
			})
			if err != nil {
				return nil, err
			}

			s.Log.Debug("[%s] state.RootModule().Outputs = %+v\n", cred.Identifier, state.RootModule().Outputs)

			if err := s.Builder.Object.Decode(state.RootModule().Outputs, resp); err != nil {
				return nil, err
			}

			s.Log.Debug("[%s] resp = %+v\n", cred.Identifier, resp)

			if err := resp.BootstrapValid(); err != nil {
				return nil, fmt.Errorf("invalid bootstrap metadata for %q: %s", cred.Identifier, err)
			}
		}

		s.Log.Debug("[%s] Bootstrap response: %+v", cred.Identifier, resp)

		if err := modelhelper.UpdateCredentialData(cred.Identifier, bson.M{
			mongodDBOperator: bson.M{
				"meta.acl":        resp.ACL,
				"meta.cidr_block": resp.CidrBlock,
				"meta.igw":        resp.IGW,
				"meta.key_pair":   resp.KeyPair,
				"meta.rtb":        resp.RTB,
				"meta.sg":         resp.SG,
				"meta.subnet":     resp.Subnet,
				"meta.vpc":        resp.VPC,
				"meta.ami":        resp.AMI,
			},
		}); err != nil {
			return nil, err
		}
	}

	return true, nil
}

func appendAWSTemplateData(template string, awsData *awsTemplateData) (string, error) {
	var data struct {
		Output   map[string]map[string]interface{} `json:"output,omitempty"`
		Resource map[string]map[string]interface{} `json:"resource,omitempty"`
		Provider map[string]map[string]interface{} `json:"provider,omitempty"`
		Variable map[string]map[string]interface{} `json:"variable,omitempty"`
	}

	if err := json.Unmarshal([]byte(template), &data); err != nil {
		return "", err
	}

	data.Variable["key_name"] = map[string]interface{}{
		"default": awsData.KeyPairName,
	}

	data.Variable["public_key"] = map[string]interface{}{
		"default": awsData.PublicKey,
	}

	data.Variable["environment_name"] = map[string]interface{}{
		"default": awsData.EnvironmentName,
	}

	out, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return "", err
	}

	return string(out), nil
}

// awsTemplateData is being used to format the bootstrap before we pass it to
// terraformer
type awsTemplateData struct {
	KeyPairName     string
	PublicKey       string
	EnvironmentName string
}

var awsBootstrap = `{
    "provider": {
        "aws": {
            "access_key": "${var.aws_access_key}",
            "secret_key": "${var.aws_secret_key}",
            "region": "${var.aws_region}"
        }
    },
    "output": {
        "vpc": {
            "value": "${aws_vpc.vpc.id}"
        },
        "cidr_block": {
            "value": "${aws_vpc.vpc.cidr_block}"
        },
        "rtb": {
            "value": "${aws_vpc.vpc.main_route_table_id}"
        },
        "acl": {
            "value": "${aws_vpc.vpc.default_network_acl_id}"
        },
        "igw": {
            "value": "${aws_internet_gateway.main_vpc_igw.id}"
        },
        "subnet": {
            "value": "${aws_subnet.main_koding_subnet.id}"
        },
        "sg": {
            "value": "${aws_security_group.allow_all.id}"
        },
        "ami": {
            "value": "${lookup(var.aws_amis, var.aws_region)}"
        },
        "key_pair": {
            "value": "${aws_key_pair.koding_key_pair.key_name}"
        }
    },
    "resource": {
        "aws_vpc": {
            "vpc": {
                "cidr_block": "${var.cidr_block}",
                "tags": {
                    "Name": "${var.environment_name}"
                }
            }
        },
        "aws_internet_gateway": {
            "main_vpc_igw": {
                "tags": {
                    "Name": "${var.environment_name}"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_subnet": {
            "main_koding_subnet": {
                "availability_zone": "${lookup(var.aws_availability_zones, var.aws_region)}",
                "cidr_block": "${var.cidr_block}",
                "map_public_ip_on_launch": true,
                "tags": {
                    "Name": "${var.environment_name}",
                    "subnet": "public"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_route_table": {
            "public": {
                "route": {
                    "cidr_block": "0.0.0.0/0",
                    "gateway_id": "${aws_internet_gateway.main_vpc_igw.id}"
                },
                "tags": {
                    "Name": "${var.environment_name}",
                    "routeTable": "koding",
                    "subnet": "public"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_route_table_association": {
            "public-1": {
                "route_table_id": "${aws_route_table.public.id}",
                "subnet_id": "${aws_subnet.main_koding_subnet.id}"
            }
        },
        "aws_security_group": {
            "allow_all": {
                "description": "Allow all inbound and outbound traffic",
                "ingress": {
                    "from_port": 0,
                    "to_port": 0,
                    "protocol": "-1",
                    "cidr_blocks": ["0.0.0.0/0"],
                    "self": true
                },
                "egress": {
                    "from_port": 0,
                    "to_port": 0,
                    "protocol": "-1",
                    "self": true,
                    "cidr_blocks": ["0.0.0.0/0"]
                },
                "name": "allow_all",
                "tags": {
                    "Name": "${var.environment_name}"
                },
                "vpc_id": "${aws_vpc.vpc.id}"
            }
        },
        "aws_key_pair": {
            "koding_key_pair": {
                "key_name": "${var.key_name}",
                "public_key": "${var.public_key}"
            }
        }
    },
    "variable": {
        "cidr_block": {
            "default": "10.0.0.0/16"
        },
        "environment_name": {
            "default": "Koding-Bootstrap"
        },
        "aws_availability_zones": {
            "default": {
                "ap-northeast-1": "ap-northeast-1b",
                "ap-southeast-1": "ap-southeast-1b",
                "ap-southeast-2": "ap-southeast-2b",
                "eu-central-1": "eu-central-1b",
                "eu-west-1": "eu-west-1b",
                "sa-east-1": "sa-east-1b",
                "us-east-1": "us-east-1b",
                "us-west-1": "us-west-1b",
                "us-west-2": "us-west-2b"
            }
        },
        "aws_amis": {
            "default": {
                "ap-northeast-1": "ami-9e5cff9e",
                "ap-southeast-1": "ami-ec7879be",
                "ap-southeast-2": "ami-2fce8b15",
                "eu-central-1": "ami-60f9c27d",
                "eu-west-1": "ami-7c4b0a0b",
                "sa-east-1": "ami-cd9518d0",
                "us-east-1": "ami-cf35f3a4",
                "us-west-1": "ami-b33dccf7",
                "us-west-2": "ami-8d5b5dbd"
            }
        }
    }
}`
<|endoftext|>


<commit_before>package main

import (
	"github.com/1lann/beacon/chat"
	"github.com/1lann/beacon/handler"
	"time"
)

type state int

const (
	stateInitializing = iota
	stateStopped
	stateOff
	stateStartShutdown
	stateShutdown
	stateSnapshot
	stateDestroy
	stateStarted
	stateStarting
	stateUnavailable
)

func (s state) String() string {
	switch s {
	case stateInitializing:
		return "Intializing"
	case stateStopped:
		return "Stopped"
	case stateOff:
		return "Off"
	case stateSnapshot:
		return "Snapshot"
	case stateShutdown:
		return "Shutdown"
	case stateDestroy:
		return "Destroy"
	case stateStarted:
		return "Started"
	case stateStarting:
		return "Starting"
	}

	return "Unknown"
}

func (s *Server) SetState(st state) {
	if st == s.State {
		return
	}

	if st != stateStarted && s.State == stateStarted || s.State == stateOff {
		handler.Handle(s.Hostnames, s.ResponseHandler)
	}

	switch st {
	case stateInitializing:
		s.ConnectMessage = "Sorry, the server is not ready take requests " +
			"yet!\nTry connecting again in a few seconds."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Yellow + "Initializing..."
		s.PingStatus.ShowConnection = false
	case stateStopped:
		s.ConnectMessage = "Sorry, the server is intentionally down.\n" +
			"Contact Chuie for more information."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Red + "Intentionally down."
		s.PingStatus.ShowConnection = false
	case stateShutdown, stateSnapshot, stateDestroy, stateStartShutdown:
		s.ConnectMessage = "Sorry, the server is currently shutting down.\n" +
			"You may start it again when it is completely powered off."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Yellow + "Shutting down..."
		s.PingStatus.ShowConnection = false
	case stateOff:
		handler.Handle(s.Hostnames, s.StartServerHandler)
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Gold + "Powered off. Connect to start."
		s.PingStatus.ShowConnection = true
	case stateStarting:
		s.ConnectMessage = "Sorry, the server is still starting up.\n" +
			"Try connecting again in a few minutes."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.LightGreen + "Starting up..."
		s.PingStatus.ShowConnection = false
	case stateUnavailable:
		s.ConnectMessage = "The server is unavailable due to an error.\n" +
			"Contact Chuie for help."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Red + "Unavailable."
		s.PingStatus.ShowConnection = false
	case stateStarted:
		s.LastConnectionTime = time.Now()
		handler.Forward(s.Hostnames, s.IPAddress+":25565")
	}

	s.State = st
}
<commit_msg>Fix states not being shown correctly<commit_after>package main

import (
	"github.com/1lann/beacon/chat"
	"github.com/1lann/beacon/handler"
	"time"
)

type state int

const (
	stateInitializing = iota
	stateStopped
	stateOff
	stateStartShutdown
	stateShutdown
	stateSnapshot
	stateDestroy
	stateStarted
	stateStarting
	stateUnavailable
)

func (s state) String() string {
	switch s {
	case stateInitializing:
		return "Intializing"
	case stateStopped:
		return "Stopped"
	case stateOff:
		return "Off"
	case stateSnapshot:
		return "Snapshot"
	case stateShutdown:
		return "Shutdown"
	case stateDestroy:
		return "Destroy"
	case stateStarted:
		return "Started"
	case stateStarting:
		return "Starting"
	}

	return "Unknown"
}

func (s *Server) SetState(st state) {
	if st == s.State && s.State != stateInitializing {
		return
	}

	if s.State == stateStarted || s.State == stateOff {
		handler.Handle(s.Hostnames, s.ResponseHandler)
	}

	switch st {
	case stateInitializing:
		s.ConnectMessage = "Sorry, the server is not ready take requests " +
			"yet!\nTry connecting again in a few seconds."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Yellow + "Initializing..."
		s.PingStatus.ShowConnection = false
	case stateStopped:
		s.ConnectMessage = "Sorry, the server is intentionally down.\n" +
			"Contact Chuie for more information."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Red + "Intentionally down."
		s.PingStatus.ShowConnection = false
	case stateShutdown, stateSnapshot, stateDestroy, stateStartShutdown:
		s.ConnectMessage = "Sorry, the server is currently shutting down.\n" +
			"You may start it again when it is completely powered off."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Yellow + "Shutting down..."
		s.PingStatus.ShowConnection = false
	case stateOff:
		handler.Handle(s.Hostnames, s.StartServerHandler)
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Gold + "Powered off. Connect to start."
		s.PingStatus.ShowConnection = true
	case stateStarting:
		s.ConnectMessage = "Sorry, the server is still starting up.\n" +
			"Try connecting again in a few minutes."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.LightGreen + "Starting up..."
		s.PingStatus.ShowConnection = false
	case stateUnavailable:
		s.ConnectMessage = "The server is unavailable due to an error.\n" +
			"Contact Chuie for help."
		s.PingStatus.Message = chat.Format(s.Messages.ServerInfoPrefix) +
			chat.Red + "Unavailable."
		s.PingStatus.ShowConnection = false
	case stateStarted:
		s.LastConnectionTime = time.Now()
		handler.Forward(s.Hostnames, s.IPAddress+":25565")
	}

	s.State = st
}
<|endoftext|>


<commit_before>// +build linux

package linux

import (
	"os/exec"
	"strings"

	"github.com/mackerelio/mackerel-agent/logging"
)

// KernelGenerator XXX
type KernelGenerator struct {
}

// Key XXX
func (g *KernelGenerator) Key() string {
	return "kernel"
}

var kernelLogger = logging.GetLogger("spec.kernel")

// Generate XXX
func (g *KernelGenerator) Generate() (interface{}, error) {
	commands := map[string][]string{
		"name":    {"uname", "-s"},
		"release": {"uname", "-r"},
		"version": {"uname", "-v"},
		"machine": {"uname", "-m"},
		"os":      {"uname", "-o"},
	}

	results := make(map[string]string)
	for key, command := range commands {
		out, err := exec.Command(command[0], command[1]).Output()
		if err != nil {
			kernelLogger.Errorf("Failed to run %s %s (skip this spec): %s", command[0], command[1], err)
			return nil, err
		}
		str := strings.TrimSpace(string(out))

		results[key] = str
	}

	return results, nil
}
<commit_msg>add linux distribution info to kernel spec<commit_after>// +build linux

package linux

import (
	"os/exec"
	"strings"

	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/shirou/gopsutil/host"
)

// KernelGenerator XXX
type KernelGenerator struct {
}

// Key XXX
func (g *KernelGenerator) Key() string {
	return "kernel"
}

var kernelLogger = logging.GetLogger("spec.kernel")

// Generate XXX
func (g *KernelGenerator) Generate() (interface{}, error) {
	commands := map[string][]string{
		"name":    {"uname", "-s"},
		"release": {"uname", "-r"},
		"version": {"uname", "-v"},
		"machine": {"uname", "-m"},
		"os":      {"uname", "-o"},
	}

	results := make(map[string]string)
	for key, command := range commands {
		out, err := exec.Command(command[0], command[1]).Output()
		if err != nil {
			kernelLogger.Errorf("Failed to run %s %s (skip this spec): %s", command[0], command[1], err)
			return nil, err
		}
		str := strings.TrimSpace(string(out))

		results[key] = str
	}

	hostInfo, err := host.Info()
	if err != nil {
		kernelLogger.Errorf("Failed to get platform information: %s", err)
	}

	if np := normalizePlatform(hostInfo.Platform); np != "" {
		platformInfo := np + " " + hostInfo.PlatformVersion

		// when platform version is empty string
		platformInfo = strings.TrimSpace(platformInfo)

		results["platform"] = platformInfo
	}

	return results, nil
}

func normalizePlatform(platform string) string {
	var normalized string

	switch platform {
	case "debian":
		normalized = "Debian"
	case "ubuntu":
		normalized = "Ubuntu"
	case "linuxmint":
		normalized = "Linux Mint"
	case "raspbian":
		normalized = "Raspbian"
	case "fedora":
		normalized = "Fedora"
	case "oracle":
		normalized = "Oracle Linux"
	case "centos":
		normalized = "CentOS"
	case "redhat":
		normalized = "Red Hat Enterprise Linux"
	case "scientific":
		normalized = "Scientific Linux"
	case "amazon":
		normalized = "Amazon Linux"
	case "xenserver":
		normalized = "XenServer"
	case "cloudlinux":
		normalized = "CloudLinux"
	case "ibm_powerkvm":
		normalized = "IBM PowerKVM"
	case "suse":
		normalized = "SUSE Linux Enterprise Server"
	case "opensuse":
		normalized = "openSUSE"
	case "gentoo":
		normalized = "Gentoo Linux"
	case "slackware":
		normalized = "Slackware"
	case "arch":
		normalized = "Arch Linux"
	case "exherbo":
		normalized = "Exherbo"
	case "alpine":
		normalized = "Alpine Linux"
	case "coreos":
		normalized = "CoreOS"
	default:
		normalized = ""
	}

	return normalized
}
<|endoftext|>


<commit_before>package yaml

/*
#cgo LDFLAGS: -lyaml
#include <yaml.h>

int load_doc(const unsigned char *input, size_t size, yaml_document_t *document) {
	yaml_parser_t parser;

	yaml_parser_initialize(&parser);
	yaml_parser_set_input_string(&parser, input, size);

	return yaml_parser_load(&parser, document);
}
*/
import "C"
import "errors"
import "unsafe"

type scalarNode struct {
	value  *C.yaml_char_t
	length C.size_t
	style  C.yaml_scalar_style_t
}

type sequenceNode struct {
	items struct {
		// actually *C.yaml_node_item_t
		start *C.int
		end   *C.int
		top   *C.int
	}
	style C.yaml_sequence_style_t
}

type mappingNode struct {
	pairs struct {
		start *C.yaml_node_pair_t
		end   *C.yaml_node_pair_t
		top   *C.yaml_node_pair_t
	}
	style C.yaml_mapping_style_t
}

func Load(data []byte) (interface{}, error) {
	if len(data) == 0 {
		return nil, nil
	}

	document := (*C.yaml_document_t)(C.malloc(C.sizeof_yaml_document_t))
	defer C.yaml_document_delete(document)

	ok := C.load_doc((*C.uchar)(unsafe.Pointer(&data[0])), C.size_t(len(data)), document)
	if ok != 1 {
		return nil, errors.New("failed to parse yaml")
	}

	root := C.yaml_document_get_root_node(document)
	return LoadNode(document, root), nil
}

func LoadNode(document *C.yaml_document_t, node *C.yaml_node_t) interface{} {
	if node == nil {
		return nil
	}

	switch node._type {
	case C.YAML_NO_NODE:
		return nil

	case C.YAML_SCALAR_NODE:
		ystr := (*scalarNode)(unsafe.Pointer(&node.data))
		rstr := C.GoBytes(unsafe.Pointer(ystr.value), C.int(ystr.length))
		//tag := C.GoString((*C.char)(unsafe.Pointer(node.tag)))
		return string(rstr)

	case C.YAML_SEQUENCE_NODE:
		yseq := (*sequenceNode)(unsafe.Pointer(&node.data))
		rseq := make([]interface{}, 0)

		start := uintptr(unsafe.Pointer(yseq.items.start))
		end := uintptr(unsafe.Pointer(yseq.items.top))
		off := unsafe.Sizeof(*yseq.items.start)

		for i := start; i < end; i += off {
			item := (*C.int)(unsafe.Pointer(i))
			itemNode := C.yaml_document_get_node(document, *item)
			elem := LoadNode(document, itemNode)
			rseq = append(rseq, elem)
		}

		return rseq

	case C.YAML_MAPPING_NODE:
		ymap := (*mappingNode)(unsafe.Pointer(&node.data))
		rmap := make(map[string]interface{}, 0)

		start := uintptr(unsafe.Pointer(ymap.pairs.start))
		end := uintptr(unsafe.Pointer(ymap.pairs.top))
		off := unsafe.Sizeof(*ymap.pairs.start)

		for i := start; i < end; i += off {
			pair := (*C.yaml_node_pair_t)(unsafe.Pointer(i))

			keyNode := C.yaml_document_get_node(document, pair.key)
			valNode := C.yaml_document_get_node(document, pair.value)
			key := LoadNode(document, keyNode)
			val := LoadNode(document, valNode)

			rmap[key.(string)] = val
		}

		return rmap

	default:
		return nil
	}
}
<commit_msg>yaml: unexport LoadNode<commit_after>package yaml

/*
#cgo LDFLAGS: -lyaml
#include <yaml.h>

int load_doc(const unsigned char *input, size_t size, yaml_document_t *document) {
	yaml_parser_t parser;

	yaml_parser_initialize(&parser);
	yaml_parser_set_input_string(&parser, input, size);

	return yaml_parser_load(&parser, document);
}
*/
import "C"
import "errors"
import "unsafe"

type scalarNode struct {
	value  *C.yaml_char_t
	length C.size_t
	style  C.yaml_scalar_style_t
}

type sequenceNode struct {
	items struct {
		// actually *C.yaml_node_item_t
		start *C.int
		end   *C.int
		top   *C.int
	}
	style C.yaml_sequence_style_t
}

type mappingNode struct {
	pairs struct {
		start *C.yaml_node_pair_t
		end   *C.yaml_node_pair_t
		top   *C.yaml_node_pair_t
	}
	style C.yaml_mapping_style_t
}

func Load(data []byte) (interface{}, error) {
	if len(data) == 0 {
		return nil, nil
	}

	document := (*C.yaml_document_t)(C.malloc(C.sizeof_yaml_document_t))
	defer C.yaml_document_delete(document)

	ok := C.load_doc((*C.uchar)(unsafe.Pointer(&data[0])), C.size_t(len(data)), document)
	if ok != 1 {
		return nil, errors.New("failed to parse yaml")
	}

	root := C.yaml_document_get_root_node(document)
	return loadNode(document, root), nil
}

func loadNode(document *C.yaml_document_t, node *C.yaml_node_t) interface{} {
	if node == nil {
		return nil
	}

	switch node._type {
	case C.YAML_NO_NODE:
		return nil

	case C.YAML_SCALAR_NODE:
		ystr := (*scalarNode)(unsafe.Pointer(&node.data))
		rstr := C.GoBytes(unsafe.Pointer(ystr.value), C.int(ystr.length))
		//tag := C.GoString((*C.char)(unsafe.Pointer(node.tag)))
		return string(rstr)

	case C.YAML_SEQUENCE_NODE:
		yseq := (*sequenceNode)(unsafe.Pointer(&node.data))
		rseq := make([]interface{}, 0)

		start := uintptr(unsafe.Pointer(yseq.items.start))
		end := uintptr(unsafe.Pointer(yseq.items.top))
		off := unsafe.Sizeof(*yseq.items.start)

		for i := start; i < end; i += off {
			item := (*C.int)(unsafe.Pointer(i))
			itemNode := C.yaml_document_get_node(document, *item)
			elem := loadNode(document, itemNode)
			rseq = append(rseq, elem)
		}

		return rseq

	case C.YAML_MAPPING_NODE:
		ymap := (*mappingNode)(unsafe.Pointer(&node.data))
		rmap := make(map[string]interface{}, 0)

		start := uintptr(unsafe.Pointer(ymap.pairs.start))
		end := uintptr(unsafe.Pointer(ymap.pairs.top))
		off := unsafe.Sizeof(*ymap.pairs.start)

		for i := start; i < end; i += off {
			pair := (*C.yaml_node_pair_t)(unsafe.Pointer(i))

			keyNode := C.yaml_document_get_node(document, pair.key)
			valNode := C.yaml_document_get_node(document, pair.value)
			key := loadNode(document, keyNode)
			val := loadNode(document, valNode)

			rmap[key.(string)] = val
		}

		return rmap

	default:
		return nil
	}
}
<|endoftext|>


<commit_before>package apirequests

import (
	"strconv"
	"sync"
	"time"
	"github.com/devicehive/devicehive-go/internal/utils"
)

func NewWSSubscriptionsMap(clients *PendingRequestsMap) *WSSubscriptionsMap {
	return &WSSubscriptionsMap{
		PendingRequestsMap: clients,
		mu:                 sync.Mutex{},
	}
}

type WSSubscriptionsMap struct {
	*PendingRequestsMap
	buffer [][]byte
	mu     sync.Mutex
}

func (s *WSSubscriptionsMap) BufferPut(b []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.buffer = append(s.buffer, b)
}

func (s *WSSubscriptionsMap) CreateSubscription(key string) *PendingRequest {
	s.mu.Lock()
	defer s.mu.Unlock()

	client := s.PendingRequestsMap.CreateSubscription(key)

	subsData, newBuffer := s.getSubscriberData(key)

	for _, b := range subsData {
		client.Data <- b
	}

	s.buffer = newBuffer

	return client
}

func (s *WSSubscriptionsMap) getSubscriberData(subsId string) (subsData [][]byte, newBuffer [][]byte) {
	for _, b := range s.buffer {
		ids, err := utils.ParseIDs(b)
		if err != nil {
			continue
		}

		dataSubsId := strconv.FormatInt(ids.Subscription, 10)

		if dataSubsId != "" && dataSubsId == subsId {
			subsData = append(subsData, b)
		} else if dataSubsId != "" && dataSubsId != subsId {
			newBuffer = append(newBuffer, b)
		}
	}

	return subsData, newBuffer
}

func (s *WSSubscriptionsMap) CleanupBufferByTimeout(timeout time.Duration) {
	for {
		time.Sleep(timeout)

		s.mu.Lock()
		l := len(s.buffer)
		if l == 0 {
			s.mu.Unlock()
			continue
		}

		m := (l - 1) / 2
		s.buffer = s.buffer[m:]
		s.mu.Unlock()
	}
}
<commit_msg>Rename method of wssubscriptionsmap<commit_after>package apirequests

import (
	"strconv"
	"sync"
	"time"
	"github.com/devicehive/devicehive-go/internal/utils"
)

func NewWSSubscriptionsMap(clients *PendingRequestsMap) *WSSubscriptionsMap {
	return &WSSubscriptionsMap{
		PendingRequestsMap: clients,
		mu:                 sync.Mutex{},
	}
}

type WSSubscriptionsMap struct {
	*PendingRequestsMap
	buffer [][]byte
	mu     sync.Mutex
}

func (s *WSSubscriptionsMap) BufferPut(b []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.buffer = append(s.buffer, b)
}

func (s *WSSubscriptionsMap) CreateSubscription(key string) *PendingRequest {
	s.mu.Lock()
	defer s.mu.Unlock()

	client := s.PendingRequestsMap.CreateSubscription(key)

	subsData, newBuffer := s.extractSubscriberData(key)

	for _, b := range subsData {
		client.Data <- b
	}

	s.buffer = newBuffer

	return client
}

func (s *WSSubscriptionsMap) extractSubscriberData(subsId string) (subsData [][]byte, newBuffer [][]byte) {
	for _, b := range s.buffer {
		ids, err := utils.ParseIDs(b)
		if err != nil {
			continue
		}

		dataSubsId := strconv.FormatInt(ids.Subscription, 10)

		if dataSubsId != "" && dataSubsId == subsId {
			subsData = append(subsData, b)
		} else if dataSubsId != "" && dataSubsId != subsId {
			newBuffer = append(newBuffer, b)
		}
	}

	return subsData, newBuffer
}

func (s *WSSubscriptionsMap) CleanupBufferByTimeout(timeout time.Duration) {
	for {
		time.Sleep(timeout)

		s.mu.Lock()
		l := len(s.buffer)
		if l == 0 {
			s.mu.Unlock()
			continue
		}

		m := (l - 1) / 2
		s.buffer = s.buffer[m:]
		s.mu.Unlock()
	}
}
<|endoftext|>


<commit_before>package yuebao

import (
    "fmt"
    //"strconv"
    "testing"
    "time"
    "sync"
)

var err error

var w sync.WaitGroup

func Test_GrabLatestData(t *testing.T) {
    fmt.Println("\nTesting GrabLatestData()...")
    fmt.Printf("Try to grab latest yuebao data.\n")
    err := GrabLatestData()
    if err != nil {
        fmt.Println(err)
        t.Error(err)
    }
}

func Concurrent_GrabLatestData(t *testing.T){
    Test_GrabLatestData(t)
    w.Done()
}

func Test_GrabHistoryData(t *testing.T) {
    fmt.Println("\nTesting GrabHistoryData()...")
    fmt.Printf("Try to grab history yuebao data.\n")
    err := GrabHistoryData()
    if err != nil {
        fmt.Println(err)
        t.Error(err)
    }
}

func Concurrent_GrabHistoryData(t *testing.T) {
    Test_GrabHistoryData(t)
    w.Done()
}

func Test_GetData(t *testing.T) {
    fmt.Println("\nTesting GetData()...")
    s := []string{"2013-04-01", "2013-10-30", "2014-10-02"}
    for _, v := range s {
        fmt.Printf("%s, data = %s\n", v, GetData(v))
    }
}

func Concurrent_GetData(t *testing.T) {
    Test_GetData(t)
    w.Done()
}

func Test_GetDataByRange(t *testing.T) {
    fmt.Println("\nTesting GetDataByRange()...")
    tm := time.Now()
    today := fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day())
    fmt.Printf("Range: 2013-05-30 - %s\n", today)
    str := GetDataByRange("2013-05-30", today)
    fmt.Println(str)
    fmt.Printf("Range: 2012-01-01 - 2013-04-30\n")
    str =  GetDataByRange("2012-01-01", "2013-04-30")
    fmt.Println(str)
}

func Concurrent_GetDataByRange(t *testing.T) {
    Test_GetDataByRange(t)
    w.Done()
}

func Test_IsDateValid(t *testing.T) {
    fmt.Println("\nTesting IsDateValid()...")
    s := []string{"2013-4-20", "July, 4, 1999", "2013-05-29", "2013-05-30", "2015-01-01"}
    tm := time.Now()
    s = append(s, fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day()))  // today
    s = append(s, fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day() + 1))  // tomorrow
    for _, v := range s {
        fmt.Printf("date = %s, valid = %v\n", v, IsDateValid(v))
    }
}

func Test_Concurrent(t *testing.T) {
    fmt.Println("\nTesting Concurrent()...")
    w.Add(5)
    go Concurrent_GrabLatestData(t)
    go Concurrent_GrabHistoryData(t)
    go Concurrent_GrabHistoryData(t)  // call 2 times to test write lock
    go Concurrent_GetData(t)
    go Concurrent_GetDataByRange(t)
    w.Wait()
}
<commit_msg>Remove unnecessary comments.<commit_after>package yuebao

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

var err error

var w sync.WaitGroup

func Test_GrabLatestData(t *testing.T) {
	fmt.Println("\nTesting GrabLatestData()...")
	fmt.Printf("Try to grab latest yuebao data.\n")
	err := GrabLatestData()
	if err != nil {
		fmt.Println(err)
		t.Error(err)
	}
}

func Concurrent_GrabLatestData(t *testing.T) {
	Test_GrabLatestData(t)
	w.Done()
}

func Test_GrabHistoryData(t *testing.T) {
	fmt.Println("\nTesting GrabHistoryData()...")
	fmt.Printf("Try to grab history yuebao data.\n")
	err := GrabHistoryData()
	if err != nil {
		fmt.Println(err)
		t.Error(err)
	}
}

func Concurrent_GrabHistoryData(t *testing.T) {
	Test_GrabHistoryData(t)
	w.Done()
}

func Test_GetData(t *testing.T) {
	fmt.Println("\nTesting GetData()...")
	s := []string{"2013-04-01", "2013-10-30", "2014-10-02"}
	for _, v := range s {
		fmt.Printf("%s, data = %s\n", v, GetData(v))
	}
}

func Concurrent_GetData(t *testing.T) {
	Test_GetData(t)
	w.Done()
}

func Test_GetDataByRange(t *testing.T) {
	fmt.Println("\nTesting GetDataByRange()...")
	tm := time.Now()
	today := fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day())
	fmt.Printf("Range: 2013-05-30 - %s\n", today)
	str := GetDataByRange("2013-05-30", today)
	fmt.Println(str)
	fmt.Printf("Range: 2012-01-01 - 2013-04-30\n")
	str = GetDataByRange("2012-01-01", "2013-04-30")
	fmt.Println(str)
}

func Concurrent_GetDataByRange(t *testing.T) {
	Test_GetDataByRange(t)
	w.Done()
}

func Test_IsDateValid(t *testing.T) {
	fmt.Println("\nTesting IsDateValid()...")
	s := []string{"2013-4-20", "July, 4, 1999", "2013-05-29", "2013-05-30", "2015-01-01"}
	tm := time.Now()
	s = append(s, fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day()))   // today
	s = append(s, fmt.Sprintf("%04d-%02d-%02d", tm.Year(), tm.Month(), tm.Day()+1)) // tomorrow
	for _, v := range s {
		fmt.Printf("date = %s, valid = %v\n", v, IsDateValid(v))
	}
}

func Test_Concurrent(t *testing.T) {
	fmt.Println("\nTesting Concurrent()...")
	w.Add(5)
	go Concurrent_GrabLatestData(t)
	go Concurrent_GrabHistoryData(t)
	go Concurrent_GrabHistoryData(t) // call 2 times to test write lock
	go Concurrent_GetData(t)
	go Concurrent_GetDataByRange(t)
	w.Wait()
}
<|endoftext|>


<commit_before>package gopikacloud

import (
	"errors"
	"fmt"
)

// ZoneRecord definition
type ZoneRecord struct {
	ID       int    `json:"id,omitempty"`
	ZoneID   int    `json:"zone,omitempty"`
	Rtype    string `json:"rtype"`
	Ipv4     string `json:"ipv4,omitempty"`
	Ipv6     string `json:"ipv6,omitempty"`
	Name     string `json:"name,omitempty"`
	Hostname string `json:"hostname,omitempty"`
	Content  string `json:"content,omitempty"`
	TTL      int    `json:"ttl,omitempty"`
}

func zoneRecordIdentifier(value interface{}) string {
	switch value := value.(type) {
	case int:
		return fmt.Sprintf("%d", value)
	case ZoneRecord:
		return fmt.Sprintf("%d", value.ID)
	}
	return ""
}

func zoneRecordPath(zone interface{}, zonerecord interface{}) string {
	str := fmt.Sprintf("zones/%s/records/", zoneIdentifier(zone))
	if zonerecord != nil {
		str += fmt.Sprintf("%s/", zoneRecordIdentifier(zonerecord))
	}
	return str
}

// ZoneRecords lists DNS zone Records of a specific zone
func (client *Client) ZoneRecords(zone interface{}) ([]ZoneRecord, error) {
	zoneRecords := []ZoneRecord{}
	if err := client.get(zoneRecordPath(zone, nil), &zoneRecords); err != nil {
		return []ZoneRecord{}, err
	}
	return zoneRecords, nil
}

// ZoneRecord fetch a single DNS zone record
func (client *Client) ZoneRecord(zone interface{}, zonerecord interface{}) (ZoneRecord, error) {
	zoneRecord := ZoneRecord{}
	if err := client.get(zoneRecordPath(zone, zonerecord), &zoneRecord); err != nil {
		return ZoneRecord{}, err
	}
	return zoneRecord, nil
}

// Delete a zone record
func (zoneRecord *ZoneRecord) Delete(client *Client) error {
	_, status, err := client.sendRequest("DELETE", zoneRecordPath(zoneRecord.ZoneID, zoneRecord.ID), nil)
	if err != nil {
		return err
	}
	if status == 204 {
		return nil
	}
	return errors.New("Failed to delete zone record")
}
<commit_msg>add priority argument<commit_after>package gopikacloud

import (
	"errors"
	"fmt"
)

// ZoneRecord definition
type ZoneRecord struct {
	ID       int    `json:"id,omitempty"`
	ZoneID   int    `json:"zone,omitempty"`
	Rtype    string `json:"rtype"`
	Ipv4     string `json:"ipv4,omitempty"`
	Ipv6     string `json:"ipv6,omitempty"`
	Name     string `json:"name,omitempty"`
	Hostname string `json:"hostname,omitempty"`
	Content  string `json:"content,omitempty"`
	TTL      int    `json:"ttl,omitempty"`
	Priority int    `json:"priority,omitempty"`
}

func zoneRecordIdentifier(value interface{}) string {
	switch value := value.(type) {
	case int:
		return fmt.Sprintf("%d", value)
	case ZoneRecord:
		return fmt.Sprintf("%d", value.ID)
	}
	return ""
}

func zoneRecordPath(zone interface{}, zonerecord interface{}) string {
	str := fmt.Sprintf("zones/%s/records/", zoneIdentifier(zone))
	if zonerecord != nil {
		str += fmt.Sprintf("%s/", zoneRecordIdentifier(zonerecord))
	}
	return str
}

// ZoneRecords lists DNS zone Records of a specific zone
func (client *Client) ZoneRecords(zone interface{}) ([]ZoneRecord, error) {
	zoneRecords := []ZoneRecord{}
	if err := client.get(zoneRecordPath(zone, nil), &zoneRecords); err != nil {
		return []ZoneRecord{}, err
	}
	return zoneRecords, nil
}

// ZoneRecord fetch a single DNS zone record
func (client *Client) ZoneRecord(zone interface{}, zonerecord interface{}) (ZoneRecord, error) {
	zoneRecord := ZoneRecord{}
	if err := client.get(zoneRecordPath(zone, zonerecord), &zoneRecord); err != nil {
		return ZoneRecord{}, err
	}
	return zoneRecord, nil
}

// Delete a zone record
func (zoneRecord *ZoneRecord) Delete(client *Client) error {
	_, status, err := client.sendRequest("DELETE", zoneRecordPath(zoneRecord.ZoneID, zoneRecord.ID), nil)
	if err != nil {
		return err
	}
	if status == 204 {
		return nil
	}
	return errors.New("Failed to delete zone record")
}
<|endoftext|>


<commit_before>// +build !lambdabinary

package sparta

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

func platformLogSysInfo(lambdaFunc string, logger *logrus.Logger) {
	// NOP
}

// RegisterCodePipelineEnvironment is part of a CodePipeline deployment
// and defines the environments available for deployment. Environments
// are defined the `environmentName`. The values defined in the
// environmentVariables are made available to each service as
// environment variables. The environment key will be transformed into
// a configuration file for a CodePipeline CloudFormation action:
// TemplateConfiguration: !Sub "TemplateSource::${environmentName}".
func RegisterCodePipelineEnvironment(environmentName string,
	environmentVariables map[string]string) error {
	if _, exists := codePipelineEnvironments[environmentName]; exists {
		return errors.Errorf("Environment (%s) has already been defined", environmentName)
	}
	codePipelineEnvironments[environmentName] = environmentVariables
	return nil
}

// NewLoggerWithFormatter returns a logger with the given formatter. If formatter
// is nil, a TTY-aware formatter is used
func NewLoggerWithFormatter(level string, formatter logrus.Formatter) (*logrus.Logger, error) {
	logger := logrus.New()
	// If there is an environment override, use that
	envLogLevel := os.Getenv(envVarLogLevel)
	if envLogLevel != "" {
		level = envLogLevel
	}

	logLevel, err := logrus.ParseLevel(level)
	if err != nil {
		return nil, err
	}

	logger.Level = logLevel
	if nil != formatter {
		logger.Formatter = formatter
	}
	logger.Out = os.Stdout
	return logger, nil
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api *API, site *S3Site) error {
	return MainEx(serviceName,
		serviceDescription,
		lambdaAWSInfos,
		api,
		site,
		nil,
		false)
}

// MainEx provides an "extended" Main that supports customizing the standard Sparta
// workflow via the `workflowHooks` parameter.
func MainEx(serviceName string,
	serviceDescription string,
	lambdaAWSInfos []*LambdaAWSInfo,
	api *API,
	site *S3Site,
	workflowHooks *WorkflowHooks,
	useCGO bool) error {
	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta v.%s powered AWS Lambda Microservice",
		serviceName,
		SpartaVersion)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
		// Save the ServiceName in case a custom command wants it
		OptionsGlobal.ServiceName = serviceName
		OptionsGlobal.ServiceDescription = serviceDescription

		validateErr := validate.Struct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}

		// Format?
		// Running in AWS?
		disableColors := OptionsGlobal.DisableColors || isRunningInAWS()
		var formatter logrus.Formatter
		switch OptionsGlobal.LogFormat {
		case "text", "txt":
			formatter = &logrus.TextFormatter{
				DisableColors: disableColors,
				FullTimestamp: OptionsGlobal.TimeStamps,
			}
		case "json":
			formatter = &logrus.JSONFormatter{}
			disableColors = true
		}
		logger, loggerErr := NewLoggerWithFormatter(OptionsGlobal.LogLevel, formatter)
		if nil != loggerErr {
			return loggerErr
		}

		// This is a NOP, but makes megacheck happy b/c it doesn't know about
		// build flags
		platformLogSysInfo("", logger)
		OptionsGlobal.Logger = logger
		welcomeMessage := fmt.Sprintf("Service: %s", serviceName)

		// Header information...
		displayPrettyHeader(headerDivider, disableColors, logger)
		// Metadata about the build...
		logger.WithFields(logrus.Fields{
			"Option":    cmd.Name(),
			"UTC":       (time.Now().UTC().Format(time.RFC3339)),
			"LinkFlags": OptionsGlobal.LinkerFlags,
		}).Info(welcomeMessage)
		logger.Info(headerDivider)

		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsProvision)

		OptionsGlobal.Logger.WithFields(logrus.Fields{
			"validateErr":      validateErr,
			"optionsProvision": optionsProvision,
		}).Debug("Provision validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := provisionBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}
			// Save the BuildID
			StampedBuildID = buildID
			return Provision(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsProvision.S3Bucket,
				useCGO,
				optionsProvision.InPlace,
				buildID,
				optionsProvision.PipelineTrigger,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				nil,
				workflowHooks,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {

			OptionsGlobal.Logger.Formatter = new(logrus.JSONFormatter)
			// Ensure the discovery service is initialized
			initializeDiscovery(OptionsGlobal.Logger)

			return Execute(serviceName,
				lambdaAWSInfos,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsDescribe)
			if nil != validateErr {
				return validateErr
			}

			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer fileWriter.Close()
			describeErr := Describe(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				fileWriter,
				workflowHooks,
				OptionsGlobal.Logger)

			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}

			return Explore(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	//////////////////////////////////////////////////////////////////////////////
	// Profile
	if nil == CommandLineOptions.Profile.RunE {
		CommandLineOptions.Profile.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsProfile)
			if nil != validateErr {
				return validateErr
			}
			return Profile(serviceName,
				serviceDescription,
				optionsProfile.S3Bucket,
				optionsProfile.Port,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Profile)

	//////////////////////////////////////////////////////////////////////////////
	// Status
	if nil == CommandLineOptions.Status.RunE {
		CommandLineOptions.Status.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsStatus)
			if nil != validateErr {
				return validateErr
			}
			return Status(serviceName,
				serviceDescription,
				optionsStatus.Redact,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Status)

	// Run it!
	executeErr := CommandLineOptions.Root.Execute()
	if executeErr != nil {
		if OptionsGlobal.Logger == nil {
			newLogger, newLoggerErr := NewLogger("info")
			if newLoggerErr != nil {
				fmt.Printf("Failed to create new logger: %v", newLoggerErr)
				newLogger = logrus.New()
			}
			OptionsGlobal.Logger = newLogger
		}
		if OptionsGlobal.Logger != nil {
			OptionsGlobal.Logger.Error(executeErr)
		} else {
			log.Printf("ERROR: %s", executeErr)
		}
	}

	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<commit_msg>Usability improvements for reporting usage<commit_after>// +build !lambdabinary

package sparta

import (
	"fmt"
	"log"
	"os"
	"time"

	"gopkg.in/go-playground/validator.v9"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

func platformLogSysInfo(lambdaFunc string, logger *logrus.Logger) {
	// NOP
}

// RegisterCodePipelineEnvironment is part of a CodePipeline deployment
// and defines the environments available for deployment. Environments
// are defined the `environmentName`. The values defined in the
// environmentVariables are made available to each service as
// environment variables. The environment key will be transformed into
// a configuration file for a CodePipeline CloudFormation action:
// TemplateConfiguration: !Sub "TemplateSource::${environmentName}".
func RegisterCodePipelineEnvironment(environmentName string,
	environmentVariables map[string]string) error {
	if _, exists := codePipelineEnvironments[environmentName]; exists {
		return errors.Errorf("Environment (%s) has already been defined", environmentName)
	}
	codePipelineEnvironments[environmentName] = environmentVariables
	return nil
}

// NewLoggerWithFormatter returns a logger with the given formatter. If formatter
// is nil, a TTY-aware formatter is used
func NewLoggerWithFormatter(level string, formatter logrus.Formatter) (*logrus.Logger, error) {
	logger := logrus.New()
	// If there is an environment override, use that
	envLogLevel := os.Getenv(envVarLogLevel)
	if envLogLevel != "" {
		level = envLogLevel
	}

	logLevel, err := logrus.ParseLevel(level)
	if err != nil {
		return nil, err
	}

	logger.Level = logLevel
	if nil != formatter {
		logger.Formatter = formatter
	}
	logger.Out = os.Stdout
	return logger, nil
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api *API, site *S3Site) error {
	return MainEx(serviceName,
		serviceDescription,
		lambdaAWSInfos,
		api,
		site,
		nil,
		false)
}

// MainEx provides an "extended" Main that supports customizing the standard Sparta
// workflow via the `workflowHooks` parameter.
func MainEx(serviceName string,
	serviceDescription string,
	lambdaAWSInfos []*LambdaAWSInfo,
	api *API,
	site *S3Site,
	workflowHooks *WorkflowHooks,
	useCGO bool) error {

	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta v.%s powered AWS Lambda Microservice",
		serviceName,
		SpartaVersion)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
		// Save the ServiceName in case a custom command wants it
		OptionsGlobal.ServiceName = serviceName
		OptionsGlobal.ServiceDescription = serviceDescription

		validateErr := validate.Struct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}

		// Format?
		// Running in AWS?
		disableColors := OptionsGlobal.DisableColors || isRunningInAWS()
		var formatter logrus.Formatter
		switch OptionsGlobal.LogFormat {
		case "text", "txt":
			formatter = &logrus.TextFormatter{
				DisableColors: disableColors,
				FullTimestamp: OptionsGlobal.TimeStamps,
			}
		case "json":
			formatter = &logrus.JSONFormatter{}
			disableColors = true
		}
		logger, loggerErr := NewLoggerWithFormatter(OptionsGlobal.LogLevel, formatter)
		if nil != loggerErr {
			return loggerErr
		}

		// This is a NOP, but makes megacheck happy b/c it doesn't know about
		// build flags
		platformLogSysInfo("", logger)
		OptionsGlobal.Logger = logger
		welcomeMessage := fmt.Sprintf("Service: %s", serviceName)

		// Header information...
		displayPrettyHeader(headerDivider, disableColors, logger)
		// Metadata about the build...
		logger.WithFields(logrus.Fields{
			"Option":    cmd.Name(),
			"UTC":       (time.Now().UTC().Format(time.RFC3339)),
			"LinkFlags": OptionsGlobal.LinkerFlags,
		}).Info(welcomeMessage)
		logger.Info(headerDivider)

		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validateErr := validate.Struct(optionsProvision)

		OptionsGlobal.Logger.WithFields(logrus.Fields{
			"validateErr":      validateErr,
			"optionsProvision": optionsProvision,
		}).Debug("Provision validation results")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			buildID, buildIDErr := provisionBuildID(optionsProvision.BuildID, OptionsGlobal.Logger)
			if nil != buildIDErr {
				return buildIDErr
			}
			// Save the BuildID
			StampedBuildID = buildID
			return Provision(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsProvision.S3Bucket,
				useCGO,
				optionsProvision.InPlace,
				buildID,
				optionsProvision.PipelineTrigger,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				nil,
				workflowHooks,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {

			OptionsGlobal.Logger.Formatter = new(logrus.JSONFormatter)
			// Ensure the discovery service is initialized
			initializeDiscovery(OptionsGlobal.Logger)

			return Execute(serviceName,
				lambdaAWSInfos,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsDescribe)
			if nil != validateErr {
				return validateErr
			}

			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer fileWriter.Close()
			describeErr := Describe(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				fileWriter,
				workflowHooks,
				OptionsGlobal.Logger)

			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}

			return Explore(serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsDescribe.S3Bucket,
				OptionsGlobal.BuildTags,
				OptionsGlobal.LinkerFlags,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	//////////////////////////////////////////////////////////////////////////////
	// Profile
	if nil == CommandLineOptions.Profile.RunE {
		CommandLineOptions.Profile.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsProfile)
			if nil != validateErr {
				return validateErr
			}
			return Profile(serviceName,
				serviceDescription,
				optionsProfile.S3Bucket,
				optionsProfile.Port,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Profile)

	//////////////////////////////////////////////////////////////////////////////
	// Status
	if nil == CommandLineOptions.Status.RunE {
		CommandLineOptions.Status.RunE = func(cmd *cobra.Command, args []string) error {
			validateErr := validate.Struct(optionsStatus)
			if nil != validateErr {
				return validateErr
			}
			return Status(serviceName,
				serviceDescription,
				optionsStatus.Redact,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Status)

	// Run it!
	executedCmd, executeErr := CommandLineOptions.Root.ExecuteC()
	if executeErr != nil {
		if OptionsGlobal.Logger == nil {
			newLogger, newLoggerErr := NewLogger("info")
			if newLoggerErr != nil {
				fmt.Printf("Failed to create new logger: %v", newLoggerErr)
				newLogger = logrus.New()
			}
			OptionsGlobal.Logger = newLogger
		}
		if OptionsGlobal.Logger != nil {
			validationErr, validationErrOk := executeErr.(validator.ValidationErrors)
			if validationErrOk {
				for _, eachError := range validationErr {
					OptionsGlobal.Logger.Error(eachError)
				}
				// Only show the usage if there were input validation errors
				if executedCmd != nil {
					executedCmd.Usage()
				}
			} else {
				OptionsGlobal.Logger.Error(executeErr)
			}
		} else {
			log.Printf("ERROR: %s", executeErr)
		}
	}

	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<|endoftext|>


<commit_before>// Package google implements logging in through Google's OpenID Connect provider.
package google

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/coreos/go-oidc"
	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/log"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/admin/directory/v1"
)

const (
	issuerURL = "https://accounts.google.com"
)

// Config holds configuration options for Google logins.
type Config struct {
	ClientID     string `json:"clientID"`
	ClientSecret string `json:"clientSecret"`
	RedirectURI  string `json:"redirectURI"`

	Scopes []string `json:"scopes"` // defaults to "profile" and "email"

	// Optional list of whitelisted domains
	// If this field is nonempty, only users from a listed domain will be allowed to log in
	HostedDomains []string `json:"hostedDomains"`

	// Optional path to service account json
	// If nonempty, and groups claim is made, will use authentication from file to
	// check groups with the admin directory api
	ServiceAccountFilePath string `json:"serviceAccountFilePath"`

	// Required if ServiceAccountFilePath
	// The email of a GSuite super user which the service account will impersonate
	// when listing groups
	AdminEmail string
}

// Open returns a connector which can be used to login users through Google.
func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, err error) {
	ctx, cancel := context.WithCancel(context.Background())

	provider, err := oidc.NewProvider(ctx, issuerURL)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("failed to get provider: %v", err)
	}

	scopes := []string{oidc.ScopeOpenID}
	if len(c.Scopes) > 0 {
		scopes = append(scopes, c.Scopes...)
	} else {
		scopes = append(scopes, "profile", "email")
	}

	clientID := c.ClientID
	return &googleConnector{
		redirectURI: c.RedirectURI,
		oauth2Config: &oauth2.Config{
			ClientID:     clientID,
			ClientSecret: c.ClientSecret,
			Endpoint:     provider.Endpoint(),
			Scopes:       scopes,
			RedirectURL:  c.RedirectURI,
		},
		verifier: provider.Verifier(
			&oidc.Config{ClientID: clientID},
		),
		logger:                 logger,
		cancel:                 cancel,
		hostedDomains:          c.HostedDomains,
		serviceAccountFilePath: c.ServiceAccountFilePath,
		adminEmail:             c.AdminEmail,
	}, nil
}

var (
	_ connector.CallbackConnector = (*googleConnector)(nil)
	_ connector.RefreshConnector  = (*googleConnector)(nil)
)

type googleConnector struct {
	redirectURI            string
	oauth2Config           *oauth2.Config
	verifier               *oidc.IDTokenVerifier
	ctx                    context.Context
	cancel                 context.CancelFunc
	logger                 log.Logger
	hostedDomains          []string
	serviceAccountFilePath string
	adminEmail             string
}

func (c *googleConnector) Close() error {
	c.cancel()
	return nil
}

func (c *googleConnector) LoginURL(s connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", callbackURL, c.redirectURI)
	}

	var opts []oauth2.AuthCodeOption
	if len(c.hostedDomains) > 0 {
		preferredDomain := c.hostedDomains[0]
		if len(c.hostedDomains) > 1 {
			preferredDomain = "*"
		}
		opts = append(opts, oauth2.SetAuthURLParam("hd", preferredDomain))
	}

	if s.OfflineAccess {
		opts = append(opts, oauth2.AccessTypeOffline, oauth2.SetAuthURLParam("prompt", "consent"))
	}
	return c.oauth2Config.AuthCodeURL(state, opts...), nil
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}

func (c *googleConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}
	token, err := c.oauth2Config.Exchange(r.Context(), q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("google: failed to get token: %v", err)
	}

	return c.createIdentity(r.Context(), identity, s, token)
}

func (c *googleConnector) Refresh(ctx context.Context, s connector.Scopes, identity connector.Identity) (connector.Identity, error) {
	t := &oauth2.Token{
		RefreshToken: string(identity.ConnectorData),
		Expiry:       time.Now().Add(-time.Hour),
	}
	token, err := c.oauth2Config.TokenSource(ctx, t).Token()
	if err != nil {
		return identity, fmt.Errorf("google: failed to get token: %v", err)
	}

	return c.createIdentity(ctx, identity, s, token)
}

func (c *googleConnector) createIdentity(ctx context.Context, identity connector.Identity, s connector.Scopes, token *oauth2.Token) (connector.Identity, error) {
	rawIDToken, ok := token.Extra("id_token").(string)
	if !ok {
		return identity, errors.New("google: no id_token in token response")
	}
	idToken, err := c.verifier.Verify(ctx, rawIDToken)
	if err != nil {
		return identity, fmt.Errorf("google: failed to verify ID Token: %v", err)
	}

	var claims struct {
		Username      string `json:"name"`
		Email         string `json:"email"`
		EmailVerified bool   `json:"email_verified"`
		HostedDomain  string `json:"hd"`
	}
	if err := idToken.Claims(&claims); err != nil {
		return identity, fmt.Errorf("oidc: failed to decode claims: %v", err)
	}

	if len(c.hostedDomains) > 0 {
		found := false
		for _, domain := range c.hostedDomains {
			if claims.HostedDomain == domain {
				found = true
				break
			}
		}

		if !found {
			return identity, fmt.Errorf("oidc: unexpected hd claim %v", claims.HostedDomain)
		}
	}

	var groups []string
	if s.Groups && c.adminEmail != "" && c.serviceAccountFilePath != "" {
		groups, err = c.getGroups(claims.Email)
		if err != nil {
			return identity, fmt.Errorf("google: could not retrieve groups: %v", err)
		}
	}

	identity = connector.Identity{
		UserID:        idToken.Subject,
		Username:      claims.Username,
		Email:         claims.Email,
		EmailVerified: claims.EmailVerified,
		ConnectorData: []byte(token.RefreshToken),
		Groups:        groups,
	}
	return identity, nil
}

// getGroups creates a connection to the admin directory service and lists
// all groups the user is a member of
func (c *googleConnector) getGroups(email string) ([]string, error) {
	srv, err := createDirectoryService(c.serviceAccountFilePath, c.adminEmail)
	if err != nil {
		return nil, fmt.Errorf("could not create directory service: %v", err)
	}

	groupsList, err := srv.Groups.List().UserKey(email).Do()
	if err != nil {
		return nil, fmt.Errorf("could not list groups: %v", err)
	}

	var userGroups []string
	for _, group := range groupsList.Groups {
		// TODO (joelspeed): Make desried group key configurable
		userGroups = append(userGroups, group.Email)
	}

	return userGroups, nil
}

// createDirectoryService loads a google service account credentials file,
// sets up super user impersonation and creates an admin client for calling
// the google admin api
func createDirectoryService(serviceAccountFilePath string, email string) (*admin.Service, error) {
	jsonCredentials, err := ioutil.ReadFile(serviceAccountFilePath)
	if err != nil {
		return nil, fmt.Errorf("error reading credentials from file: %v", err)
	}

	config, err := google.JWTConfigFromJSON(jsonCredentials, admin.AdminDirectoryGroupReadonlyScope)
	if err != nil {
		return nil, fmt.Errorf("unable to parse client secret file to config: %v", err)
	}

	// Impersonate an admin. This is mandatory for the admin APIs.
	config.Subject = email

	ctx := context.Background()
	client := config.Client(ctx)

	srv, err := admin.New(client)
	if err != nil {
		return nil, fmt.Errorf("unable to create directory service %v", err)
	}
	return srv, nil
}
<commit_msg>Make directory service during init<commit_after>// Package google implements logging in through Google's OpenID Connect provider.
package google

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/coreos/go-oidc"
	"golang.org/x/oauth2"

	"github.com/dexidp/dex/connector"
	"github.com/dexidp/dex/pkg/log"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/admin/directory/v1"
)

const (
	issuerURL = "https://accounts.google.com"
)

// Config holds configuration options for Google logins.
type Config struct {
	ClientID     string `json:"clientID"`
	ClientSecret string `json:"clientSecret"`
	RedirectURI  string `json:"redirectURI"`

	Scopes []string `json:"scopes"` // defaults to "profile" and "email"

	// Optional list of whitelisted domains
	// If this field is nonempty, only users from a listed domain will be allowed to log in
	HostedDomains []string `json:"hostedDomains"`

	// Optional path to service account json
	// If nonempty, and groups claim is made, will use authentication from file to
	// check groups with the admin directory api
	ServiceAccountFilePath string `json:"serviceAccountFilePath"`

	// Required if ServiceAccountFilePath
	// The email of a GSuite super user which the service account will impersonate
	// when listing groups
	AdminEmail string
}

// Open returns a connector which can be used to login users through Google.
func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, err error) {
	ctx, cancel := context.WithCancel(context.Background())

	provider, err := oidc.NewProvider(ctx, issuerURL)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("failed to get provider: %v", err)
	}

	scopes := []string{oidc.ScopeOpenID}
	if len(c.Scopes) > 0 {
		scopes = append(scopes, c.Scopes...)
	} else {
		scopes = append(scopes, "profile", "email")
	}

	srv, err := createDirectoryService(c.ServiceAccountFilePath, c.AdminEmail)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("could not create directory service: %v", err)
	}

	clientID := c.ClientID
	return &googleConnector{
		redirectURI: c.RedirectURI,
		oauth2Config: &oauth2.Config{
			ClientID:     clientID,
			ClientSecret: c.ClientSecret,
			Endpoint:     provider.Endpoint(),
			Scopes:       scopes,
			RedirectURL:  c.RedirectURI,
		},
		verifier: provider.Verifier(
			&oidc.Config{ClientID: clientID},
		),
		logger:                 logger,
		cancel:                 cancel,
		hostedDomains:          c.HostedDomains,
		serviceAccountFilePath: c.ServiceAccountFilePath,
		adminEmail:             c.AdminEmail,
		adminSrv:               srv,
	}, nil
}

var (
	_ connector.CallbackConnector = (*googleConnector)(nil)
	_ connector.RefreshConnector  = (*googleConnector)(nil)
)

type googleConnector struct {
	redirectURI            string
	oauth2Config           *oauth2.Config
	verifier               *oidc.IDTokenVerifier
	ctx                    context.Context
	cancel                 context.CancelFunc
	logger                 log.Logger
	hostedDomains          []string
	serviceAccountFilePath string
	adminEmail             string
	adminSrv               *admin.Service
}

func (c *googleConnector) Close() error {
	c.cancel()
	return nil
}

func (c *googleConnector) LoginURL(s connector.Scopes, callbackURL, state string) (string, error) {
	if c.redirectURI != callbackURL {
		return "", fmt.Errorf("expected callback URL %q did not match the URL in the config %q", callbackURL, c.redirectURI)
	}

	var opts []oauth2.AuthCodeOption
	if len(c.hostedDomains) > 0 {
		preferredDomain := c.hostedDomains[0]
		if len(c.hostedDomains) > 1 {
			preferredDomain = "*"
		}
		opts = append(opts, oauth2.SetAuthURLParam("hd", preferredDomain))
	}

	if s.OfflineAccess {
		opts = append(opts, oauth2.AccessTypeOffline, oauth2.SetAuthURLParam("prompt", "consent"))
	}
	return c.oauth2Config.AuthCodeURL(state, opts...), nil
}

type oauth2Error struct {
	error            string
	errorDescription string
}

func (e *oauth2Error) Error() string {
	if e.errorDescription == "" {
		return e.error
	}
	return e.error + ": " + e.errorDescription
}

func (c *googleConnector) HandleCallback(s connector.Scopes, r *http.Request) (identity connector.Identity, err error) {
	q := r.URL.Query()
	if errType := q.Get("error"); errType != "" {
		return identity, &oauth2Error{errType, q.Get("error_description")}
	}
	token, err := c.oauth2Config.Exchange(r.Context(), q.Get("code"))
	if err != nil {
		return identity, fmt.Errorf("google: failed to get token: %v", err)
	}

	return c.createIdentity(r.Context(), identity, s, token)
}

func (c *googleConnector) Refresh(ctx context.Context, s connector.Scopes, identity connector.Identity) (connector.Identity, error) {
	t := &oauth2.Token{
		RefreshToken: string(identity.ConnectorData),
		Expiry:       time.Now().Add(-time.Hour),
	}
	token, err := c.oauth2Config.TokenSource(ctx, t).Token()
	if err != nil {
		return identity, fmt.Errorf("google: failed to get token: %v", err)
	}

	return c.createIdentity(ctx, identity, s, token)
}

func (c *googleConnector) createIdentity(ctx context.Context, identity connector.Identity, s connector.Scopes, token *oauth2.Token) (connector.Identity, error) {
	rawIDToken, ok := token.Extra("id_token").(string)
	if !ok {
		return identity, errors.New("google: no id_token in token response")
	}
	idToken, err := c.verifier.Verify(ctx, rawIDToken)
	if err != nil {
		return identity, fmt.Errorf("google: failed to verify ID Token: %v", err)
	}

	var claims struct {
		Username      string `json:"name"`
		Email         string `json:"email"`
		EmailVerified bool   `json:"email_verified"`
		HostedDomain  string `json:"hd"`
	}
	if err := idToken.Claims(&claims); err != nil {
		return identity, fmt.Errorf("oidc: failed to decode claims: %v", err)
	}

	if len(c.hostedDomains) > 0 {
		found := false
		for _, domain := range c.hostedDomains {
			if claims.HostedDomain == domain {
				found = true
				break
			}
		}

		if !found {
			return identity, fmt.Errorf("oidc: unexpected hd claim %v", claims.HostedDomain)
		}
	}

	var groups []string
	if s.Groups && c.adminEmail != "" && c.serviceAccountFilePath != "" {
		groups, err = c.getGroups(claims.Email)
		if err != nil {
			return identity, fmt.Errorf("google: could not retrieve groups: %v", err)
		}
	}

	identity = connector.Identity{
		UserID:        idToken.Subject,
		Username:      claims.Username,
		Email:         claims.Email,
		EmailVerified: claims.EmailVerified,
		ConnectorData: []byte(token.RefreshToken),
		Groups:        groups,
	}
	return identity, nil
}

// getGroups creates a connection to the admin directory service and lists
// all groups the user is a member of
func (c *googleConnector) getGroups(email string) ([]string, error) {
	groupsList, err := c.adminSrv.Groups.List().UserKey(email).Do()
	if err != nil {
		return nil, fmt.Errorf("could not list groups: %v", err)
	}

	var userGroups []string
	for _, group := range groupsList.Groups {
		// TODO (joelspeed): Make desried group key configurable
		userGroups = append(userGroups, group.Email)
	}

	return userGroups, nil
}

// createDirectoryService loads a google service account credentials file,
// sets up super user impersonation and creates an admin client for calling
// the google admin api
func createDirectoryService(serviceAccountFilePath string, email string) (*admin.Service, error) {
	jsonCredentials, err := ioutil.ReadFile(serviceAccountFilePath)
	if err != nil {
		return nil, fmt.Errorf("error reading credentials from file: %v", err)
	}

	config, err := google.JWTConfigFromJSON(jsonCredentials, admin.AdminDirectoryGroupReadonlyScope)
	if err != nil {
		return nil, fmt.Errorf("unable to parse client secret file to config: %v", err)
	}

	// Impersonate an admin. This is mandatory for the admin APIs.
	config.Subject = email

	ctx := context.Background()
	client := config.Client(ctx)

	srv, err := admin.New(client)
	if err != nil {
		return nil, fmt.Errorf("unable to create directory service %v", err)
	}
	return srv, nil
}
<|endoftext|>


<commit_before>// Package chow implements Chow et al.'s white-box AES construction. There are two separate attacks on this white-box
// construction implemented in cryptanalysis/chow and cryptanalysis/chow2.
//
// Broadly, it works by taking a normal AES key and stretching it out and adding randomness to encryption computations.
// This randomness eventually cancels out and gives a correct AES encryption of a plaintext, without leaking the key. We
// can also modify the white-box key such that the function isn't exactly ct = AES(pt), but a masked or encoded function
// like ct' = Q(AES(P(pt))), where Q and P are randomly chosen affine transformations.
//
// We start by generating a white-boxed key:
//   opts := common.IndependentMasks{common.RandomMask, common.RandomMask} // Random input and output masks.
//   constr, input, output := chow.GenerateEncryptionKeys(key, seed, opts) // key is the AES key, seed is the seed for the RNG.
// which we can use to encrypt data just like a normal AES cipher:
//   constr.Encrypt(dst, src)
//
// AES white-boxes are asymmetric, meaning you have to choose whether to generate encryption or decryption keys because
// encryption keys can't be used for decryption and vice versa. Above we showed encryption; decryption is similar:
//   opts := common.IndependentMasks{common.RandomMask, common.RandomMask}
//   constr, input, output := chow.GenerateDecryptionKeys(key, seed, opts)
//   ...
//   constr.Decrypt(dst, src)
//
// There are two types of mask: common.RandomMask and common.IdentityMask. RandomMask is a random linear transformation
// and IdentityMask is the identity transformation.
//
// There are three types of ways to attach masks to the white-box: common.IndependentMasks, common.SameMasks, and
// common.MatchingMasks. IndependentMasks specifies and chooses the input and output masks independently of each other.
// SameMasks chooses a mask of the specified type and puts the same one on the input and output. MatchingMasks chooses a
// random mask for the input and puts the inverse mask on the output.
//
// The Encrypt benchmark takes about 47,000ns per Encrypt call. For context, with hardware implementations of AES, an
// Encrypt call can take as little at 30ns per Encrypt.  Heavily optimized software implementations take about 170ns per
// Encrypt.  White-Boxing AES with Chow's construction seems to make it approximately 2 to 3 orders of magnitude slower.
//
// "White-Box Cryptography and an AES Implementation" by Stanley Chow, Philip Eisen, Harold Johnson, and Paul C. Van
// Oorschot, http://link.springer.com/chapter/10.1007%2F3-540-36492-7_17?LI=true
//
// "A Tutorial on White-Box AES" by James A. Muir, https://eprint.iacr.org/2013/104.pdf
package chow

import (
	"github.com/OpenWhiteBox/primitives/table"

	"github.com/OpenWhiteBox/AES/constructions/common"
)

type Construction struct {
	InputMask      [16]table.Block // [round]
	InputXORTables common.NibbleXORTables

	TBoxTyiTable [9][16]table.Word      // [round][position]
	HighXORTable [9][32][3]table.Nibble // [round][nibble-wise position][gate number]

	MBInverseTable [9][16]table.Word      // [round][position]
	LowXORTable    [9][32][3]table.Nibble // [round][nibble-wise position][gate number]

	TBoxOutputMask  [16]table.Block // [position]
	OutputXORTables common.NibbleXORTables
}

// BlockSize returns the block size of AES. (Necessary to implement cipher.Block.)
func (constr Construction) BlockSize() int { return 16 }

// Encrypt encrypts the first block in src into dst. Dst and src may point at the same memory.
func (constr Construction) Encrypt(dst, src []byte) {
	constr.crypt(dst, src, constr.ShiftRows)
}

// Decrypt decrypts the first block in src into dst. Dst and src may point at the same memory.
func (constr Construction) Decrypt(dst, src []byte) {
	constr.crypt(dst, src, constr.UnShiftRows)
}

// crypt pushes the first block in src through the lookup tables (which may compute encryption or decryption) and writes
// the result to dst. shift is the permutation to apply to the state matrix before each round.
func (constr Construction) crypt(dst, src []byte, shift func([]byte)) {
	copy(dst, src[:constr.BlockSize()])

	// Remove input encoding.
	stretched := constr.ExpandBlock(constr.InputMask, dst)
	constr.InputXORTables.SquashBlocks(stretched, dst)

	for round := 0; round < 9; round++ {
		shift(dst)

		// Apply the T-Boxes and Tyi Tables to each column of the state matrix.
		for pos := 0; pos < 16; pos += 4 {
			stretched := constr.ExpandWord(constr.TBoxTyiTable[round][pos:pos+4], dst[pos:pos+4])
			constr.SquashWords(constr.HighXORTable[round][2*pos:2*pos+8], stretched, dst[pos:pos+4])

			stretched = constr.ExpandWord(constr.MBInverseTable[round][pos:pos+4], dst[pos:pos+4])
			constr.SquashWords(constr.LowXORTable[round][2*pos:2*pos+8], stretched, dst[pos:pos+4])
		}
	}

	shift(dst)

	// Apply the final T-Box transformation and add the output encoding.
	stretched = constr.ExpandBlock(constr.TBoxOutputMask, dst)
	constr.OutputXORTables.SquashBlocks(stretched, dst)
}

// ShiftRows permutes the bytes of the first block of block, according to AES' ShiftRows operation.
func (constr *Construction) ShiftRows(block []byte) {
	copy(block, []byte{
		block[0], block[5], block[10], block[15], block[4], block[9], block[14], block[3], block[8], block[13], block[2],
		block[7], block[12], block[1], block[6], block[11],
	})
}

// UnShiftRows permutes the bytes of the first block of block, according to the inverse of AES's ShiftRows operation.
func (constr *Construction) UnShiftRows(block []byte) {
	copy(block, []byte{
		block[0], block[13], block[10], block[7], block[4], block[1], block[14], block[11], block[8], block[5], block[2],
		block[15], block[12], block[9], block[6], block[3],
	})
}

// ExpandWord expands one word of the state matrix with the T-Boxes composed with Tyi Tables.
func (constr *Construction) ExpandWord(tboxtyi []table.Word, word []byte) [4][4]byte {
	return [4][4]byte{tboxtyi[0].Get(word[0]), tboxtyi[1].Get(word[1]), tboxtyi[2].Get(word[2]), tboxtyi[3].Get(word[3])}
}

// SquashWords squashes an expanded word back into one word with 3 pairwise XORs (calc'd one nibble at a time):
//   (((a ^ b) ^ c) ^ d)
func (constr *Construction) SquashWords(xorTable [][3]table.Nibble, words [4][4]byte, dst []byte) {
	copy(dst, words[0][:])

	for i := 1; i < 4; i++ {
		for pos := 0; pos < 4; pos++ {
			aPartial := dst[pos]&0xf0 | (words[i][pos]&0xf0)>>4
			bPartial := (dst[pos]&0x0f)<<4 | words[i][pos]&0x0f

			dst[pos] = xorTable[2*pos+0][i-1].Get(aPartial)<<4 | xorTable[2*pos+1][i-1].Get(bPartial)
		}
	}
}

// ExpandBlock expands the entire state matrix into sixteen blocks.
func (constr *Construction) ExpandBlock(mask [16]table.Block, block []byte) (out [16][16]byte) {
	for i := 0; i < 16; i++ {
		out[i] = mask[i].Get(block[i])
	}

	return
}
<commit_msg>Rephrase overview of constructions/chow.<commit_after>// Package chow implements Chow et al.'s white-box AES construction. There are two separate attacks on this white-box
// construction implemented in the cryptanalysis/chow and cryptanalysis/chow2 packages.
//
// Broadly, it works by taking a normal AES key and converting the encryption algorithm into a series of table lookups.
// The table lookups are then randomized such that this randomness eventually cancels out and gives a correct AES
// encryption of a plaintext, without leaking the key. We can also modify the white-box key such that the function isn't
// exactly ct = AES(pt), but a masked or encoded function like ct' = Q(AES(P(pt))), where Q and P are randomly chosen
// affine transformations.
//
// We start by generating a white-boxed key:
//   opts := common.IndependentMasks{common.RandomMask, common.RandomMask} // Random input and output masks.
//   constr, input, output := chow.GenerateEncryptionKeys(key, seed, opts) // key is the AES key, seed is the seed for the RNG.
// which we can use to encrypt data just like a normal AES cipher:
//   constr.Encrypt(dst, src)
//
// AES white-boxes are asymmetric, meaning you have to choose whether to generate encryption or decryption keys because
// encryption keys can't be used for decryption and vice versa. Above we showed encryption; decryption is similar:
//   opts := common.IndependentMasks{common.RandomMask, common.RandomMask}
//   constr, input, output := chow.GenerateDecryptionKeys(key, seed, opts)
//   ...
//   constr.Decrypt(dst, src)
//
// There are two types of mask: common.RandomMask and common.IdentityMask. RandomMask is a random linear transformation
// and IdentityMask is the identity transformation.
//
// There are three types of ways to attach masks to the white-box: common.IndependentMasks, common.SameMasks, and
// common.MatchingMasks. IndependentMasks specifies and chooses the input and output masks independently of each other.
// SameMasks chooses a mask of the specified type and puts the same one on the input and output. MatchingMasks chooses a
// random mask for the input and puts the inverse mask on the output.
//
// The Encrypt benchmark takes about 47,000ns per Encrypt call. For context, with hardware implementations of AES, an
// Encrypt call can take as little at 30ns per Encrypt.  Heavily optimized software implementations take about 170ns per
// Encrypt.  White-Boxing AES with Chow's construction seems to make it approximately 2 to 3 orders of magnitude slower.
//
// "White-Box Cryptography and an AES Implementation" by Stanley Chow, Philip Eisen, Harold Johnson, and Paul C. Van
// Oorschot, http://link.springer.com/chapter/10.1007%2F3-540-36492-7_17?LI=true
//
// "A Tutorial on White-Box AES" by James A. Muir, https://eprint.iacr.org/2013/104.pdf
package chow

import (
	"github.com/OpenWhiteBox/primitives/table"

	"github.com/OpenWhiteBox/AES/constructions/common"
)

type Construction struct {
	InputMask      [16]table.Block // [round]
	InputXORTables common.NibbleXORTables

	TBoxTyiTable [9][16]table.Word      // [round][position]
	HighXORTable [9][32][3]table.Nibble // [round][nibble-wise position][gate number]

	MBInverseTable [9][16]table.Word      // [round][position]
	LowXORTable    [9][32][3]table.Nibble // [round][nibble-wise position][gate number]

	TBoxOutputMask  [16]table.Block // [position]
	OutputXORTables common.NibbleXORTables
}

// BlockSize returns the block size of AES. (Necessary to implement cipher.Block.)
func (constr Construction) BlockSize() int { return 16 }

// Encrypt encrypts the first block in src into dst. Dst and src may point at the same memory.
func (constr Construction) Encrypt(dst, src []byte) {
	constr.crypt(dst, src, constr.ShiftRows)
}

// Decrypt decrypts the first block in src into dst. Dst and src may point at the same memory.
func (constr Construction) Decrypt(dst, src []byte) {
	constr.crypt(dst, src, constr.UnShiftRows)
}

// crypt pushes the first block in src through the lookup tables (which may compute encryption or decryption) and writes
// the result to dst. shift is the permutation to apply to the state matrix before each round.
func (constr Construction) crypt(dst, src []byte, shift func([]byte)) {
	copy(dst, src[:constr.BlockSize()])

	// Remove input encoding.
	stretched := constr.ExpandBlock(constr.InputMask, dst)
	constr.InputXORTables.SquashBlocks(stretched, dst)

	for round := 0; round < 9; round++ {
		shift(dst)

		// Apply the T-Boxes and Tyi Tables to each column of the state matrix.
		for pos := 0; pos < 16; pos += 4 {
			stretched := constr.ExpandWord(constr.TBoxTyiTable[round][pos:pos+4], dst[pos:pos+4])
			constr.SquashWords(constr.HighXORTable[round][2*pos:2*pos+8], stretched, dst[pos:pos+4])

			stretched = constr.ExpandWord(constr.MBInverseTable[round][pos:pos+4], dst[pos:pos+4])
			constr.SquashWords(constr.LowXORTable[round][2*pos:2*pos+8], stretched, dst[pos:pos+4])
		}
	}

	shift(dst)

	// Apply the final T-Box transformation and add the output encoding.
	stretched = constr.ExpandBlock(constr.TBoxOutputMask, dst)
	constr.OutputXORTables.SquashBlocks(stretched, dst)
}

// ShiftRows permutes the bytes of the first block of block, according to AES' ShiftRows operation.
func (constr *Construction) ShiftRows(block []byte) {
	copy(block, []byte{
		block[0], block[5], block[10], block[15], block[4], block[9], block[14], block[3], block[8], block[13], block[2],
		block[7], block[12], block[1], block[6], block[11],
	})
}

// UnShiftRows permutes the bytes of the first block of block, according to the inverse of AES's ShiftRows operation.
func (constr *Construction) UnShiftRows(block []byte) {
	copy(block, []byte{
		block[0], block[13], block[10], block[7], block[4], block[1], block[14], block[11], block[8], block[5], block[2],
		block[15], block[12], block[9], block[6], block[3],
	})
}

// ExpandWord expands one word of the state matrix with the T-Boxes composed with Tyi Tables.
func (constr *Construction) ExpandWord(tboxtyi []table.Word, word []byte) [4][4]byte {
	return [4][4]byte{tboxtyi[0].Get(word[0]), tboxtyi[1].Get(word[1]), tboxtyi[2].Get(word[2]), tboxtyi[3].Get(word[3])}
}

// SquashWords squashes an expanded word back into one word with 3 pairwise XORs (calc'd one nibble at a time):
//   (((a ^ b) ^ c) ^ d)
func (constr *Construction) SquashWords(xorTable [][3]table.Nibble, words [4][4]byte, dst []byte) {
	copy(dst, words[0][:])

	for i := 1; i < 4; i++ {
		for pos := 0; pos < 4; pos++ {
			aPartial := dst[pos]&0xf0 | (words[i][pos]&0xf0)>>4
			bPartial := (dst[pos]&0x0f)<<4 | words[i][pos]&0x0f

			dst[pos] = xorTable[2*pos+0][i-1].Get(aPartial)<<4 | xorTable[2*pos+1][i-1].Get(bPartial)
		}
	}
}

// ExpandBlock expands the entire state matrix into sixteen blocks.
func (constr *Construction) ExpandBlock(mask [16]table.Block, block []byte) (out [16][16]byte) {
	for i := 0; i < 16; i++ {
		out[i] = mask[i].Get(block[i])
	}

	return
}
<|endoftext|>