<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
)

type Config struct {
	Consul	     *string          `json:"consul"`
	Services     []*ServiceConfig `json:"services"`
	Backends     []*BackendConfig `json:"backends"`
}

type ServiceConfig struct {
	Id               string
	Name             string `json:"name"`
	Poll             int    `json:"poll"` // time in seconds
	HealthCheckExec  string `json:"health"`
	Port             int    `json:"port"`
	TTL              int    `json:"ttl"`
	IsPublic         bool   `json:"publicIp"` // will default to false
	discoveryService DiscoveryService
	healthArgs       []string
	ipAddress        string
}

type BackendConfig struct {
	Name             string `json:"name"`
	Poll             int    `json:"poll"` // time in seconds
	OnChangeExec     string `json:"onChange"`
	discoveryService DiscoveryService
	onChangeArgs     []string
	lastState        interface{}
}

type Pollable interface {
	PollTime() int
}

func (b BackendConfig) PollTime() int {
	return b.Poll
}
func (b *BackendConfig) CheckForUpstreamChanges() bool {
	return b.discoveryService.CheckForUpstreamChanges(b)
}

func (s ServiceConfig) PollTime() int {
	return s.Poll
}
func (s *ServiceConfig) WriteHealthCheck() {
	s.discoveryService.WriteHealthCheck(s)
}

func loadConfig() *Config {

	var configFlag string
	var discovery DiscoveryService
	discoveryCount := 0
	flag.StringVar(&configFlag, "config", "", "JSON config or file:// path to JSON config file.")
	flag.Parse()
	if configFlag == "" {
		configFlag = os.Getenv("CONTAINERBUDDY")
	}

	config := parseConfig(configFlag)

	for _, discoveryBackend := range []string{"Consul"} {
		switch discoveryBackend {
		case "Consul":
			if config.Consul != nil {
				discovery = NewConsulConfig(*config.Consul)
				discoveryCount += 1
			}
		}
	}

	if discoveryCount == 0 {
		log.Fatal("No discovery backend defined")
	} else if discoveryCount > 1 {
		log.Fatal("More than one discovery backend defined")
	}

	for _, backend := range config.Backends {
		backend.discoveryService = discovery
		backend.onChangeArgs = strings.Split(backend.OnChangeExec, " ")
	}

	hostname, _ := os.Hostname()
	for _, service := range config.Services {
		service.Id = fmt.Sprintf("%s-%s", service.Name, hostname)
		service.discoveryService = discovery
		service.healthArgs = strings.Split(service.HealthCheckExec, " ")
		service.ipAddress = getIp(service.IsPublic)
	}

	return config
}

func parseConfig(configFlag string) *Config {
	if configFlag == "" {
		log.Fatal("-config flag is required.")
	}

	var data []byte
	if strings.HasPrefix(configFlag, "file://") {
		var err error
		if data, err = ioutil.ReadFile(strings.SplitAfter(configFlag, "file://")[1]); err != nil {
			log.Fatalf("Could not read config file: %s", err)
		}
	} else {
		data = []byte(configFlag)
	}

	config := &Config{}
	if err := json.Unmarshal(data, &config); err != nil {
		log.Fatalf("Could not parse configuration: %s", err)
	}

	return config
}

// determine the IP address of the container
func getIp(usePublic bool) string {
	interfaces, _ := net.Interfaces()
	var ips []net.IP
	_, loopback, _ := net.ParseCIDR("127.0.0.0/8")
	for _, intf := range interfaces {
		ipAddrs, _ := intf.Addrs()
		// We're assuming each interface has one IP here because neither Docker
		// nor Triton sets up IP aliasing.
		ipAddr, _, _ := net.ParseCIDR(ipAddrs[0].String())
		if !loopback.Contains(ipAddr) {
			ips = append(ips, ipAddr)
		}
	}
	var ip string
	for _, ipAddr := range ips {
		isPublic := isPublicIp(ipAddr)
		if isPublic && usePublic {
			ip = ipAddr.String()
			break
		} else if !isPublic && !usePublic {
			ip = ipAddr.String()
			break
		}
	}
	return ip
}

// parse an IPv4 address and return true if it's a public IP
func isPublicIp(ip net.IP) bool {
	_, c, _ := net.ParseCIDR("192.168.0.0/16")
	_, b, _ := net.ParseCIDR("172.16.0.0/12")
	_, a, _ := net.ParseCIDR("10.0.0.0/8")

	var privateNetworks = []*net.IPNet{c, b, a}
	for _, network := range privateNetworks {
		if network.Contains(ip) {
			return false
		}
	}
	return true
}
<commit_msg>Use `omitempty' and strings on backend config<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
)

type Config struct {
	Consul	     string          `json:"consul,omitempty"`
	Services     []*ServiceConfig `json:"services"`
	Backends     []*BackendConfig `json:"backends"`
}

type ServiceConfig struct {
	Id               string
	Name             string `json:"name"`
	Poll             int    `json:"poll"` // time in seconds
	HealthCheckExec  string `json:"health"`
	Port             int    `json:"port"`
	TTL              int    `json:"ttl"`
	IsPublic         bool   `json:"publicIp"` // will default to false
	discoveryService DiscoveryService
	healthArgs       []string
	ipAddress        string
}

type BackendConfig struct {
	Name             string `json:"name"`
	Poll             int    `json:"poll"` // time in seconds
	OnChangeExec     string `json:"onChange"`
	discoveryService DiscoveryService
	onChangeArgs     []string
	lastState        interface{}
}

type Pollable interface {
	PollTime() int
}

func (b BackendConfig) PollTime() int {
	return b.Poll
}
func (b *BackendConfig) CheckForUpstreamChanges() bool {
	return b.discoveryService.CheckForUpstreamChanges(b)
}

func (s ServiceConfig) PollTime() int {
	return s.Poll
}
func (s *ServiceConfig) WriteHealthCheck() {
	s.discoveryService.WriteHealthCheck(s)
}

func loadConfig() *Config {

	var configFlag string
	var discovery DiscoveryService
	discoveryCount := 0
	flag.StringVar(&configFlag, "config", "", "JSON config or file:// path to JSON config file.")
	flag.Parse()
	if configFlag == "" {
		configFlag = os.Getenv("CONTAINERBUDDY")
	}

	config := parseConfig(configFlag)

	for _, discoveryBackend := range []string{"Consul"} {
		switch discoveryBackend {
		case "Consul":
			if config.Consul != "" {
				discovery = NewConsulConfig(config.Consul)
				discoveryCount += 1
			}
		}
	}

	if discoveryCount == 0 {
		log.Fatal("No discovery backend defined")
	} else if discoveryCount > 1 {
		log.Fatal("More than one discovery backend defined")
	}

	for _, backend := range config.Backends {
		backend.discoveryService = discovery
		backend.onChangeArgs = strings.Split(backend.OnChangeExec, " ")
	}

	hostname, _ := os.Hostname()
	for _, service := range config.Services {
		service.Id = fmt.Sprintf("%s-%s", service.Name, hostname)
		service.discoveryService = discovery
		service.healthArgs = strings.Split(service.HealthCheckExec, " ")
		service.ipAddress = getIp(service.IsPublic)
	}

	return config
}

func parseConfig(configFlag string) *Config {
	if configFlag == "" {
		log.Fatal("-config flag is required.")
	}

	var data []byte
	if strings.HasPrefix(configFlag, "file://") {
		var err error
		if data, err = ioutil.ReadFile(strings.SplitAfter(configFlag, "file://")[1]); err != nil {
			log.Fatalf("Could not read config file: %s", err)
		}
	} else {
		data = []byte(configFlag)
	}

	config := &Config{}
	if err := json.Unmarshal(data, &config); err != nil {
		log.Fatalf("Could not parse configuration: %s", err)
	}

	return config
}

// determine the IP address of the container
func getIp(usePublic bool) string {
	interfaces, _ := net.Interfaces()
	var ips []net.IP
	_, loopback, _ := net.ParseCIDR("127.0.0.0/8")
	for _, intf := range interfaces {
		ipAddrs, _ := intf.Addrs()
		// We're assuming each interface has one IP here because neither Docker
		// nor Triton sets up IP aliasing.
		ipAddr, _, _ := net.ParseCIDR(ipAddrs[0].String())
		if !loopback.Contains(ipAddr) {
			ips = append(ips, ipAddr)
		}
	}
	var ip string
	for _, ipAddr := range ips {
		isPublic := isPublicIp(ipAddr)
		if isPublic && usePublic {
			ip = ipAddr.String()
			break
		} else if !isPublic && !usePublic {
			ip = ipAddr.String()
			break
		}
	}
	return ip
}

// parse an IPv4 address and return true if it's a public IP
func isPublicIp(ip net.IP) bool {
	_, c, _ := net.ParseCIDR("192.168.0.0/16")
	_, b, _ := net.ParseCIDR("172.16.0.0/12")
	_, a, _ := net.ParseCIDR("10.0.0.0/8")

	var privateNetworks = []*net.IPNet{c, b, a}
	for _, network := range privateNetworks {
		if network.Contains(ip) {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package provision

import (
    "bytes"
    "fmt"
	"io/ioutil"
	"path"
    "strings"
    "text/template"

    "github.com/docker/machine/libmachine/auth"
    "github.com/docker/machine/libmachine/cert"
    "github.com/docker/machine/libmachine/kubernetes"
    "github.com/docker/machine/libmachine/log"
)

func xferCert(p Provisioner, certPath string, targetPath string) error {
	certXferCmd := "printf '%%s' '%s' | sudo tee %s"

	certContents, err := ioutil.ReadFile(certPath)
	if err != nil {
		return err
	}

    if _, err := p.SSHCommand(fmt.Sprintf("mkdir -p %s", targetPath)); err != nil {
      return err
    }

    /*
     * TODO: Until we start dynamically generating the configuration file,
     * these must have a known naming convention on the machine.
     */
     _, certFile := path.Split(certPath)
     certFile = strings.Split(certFile, "_")[1]

	if _, err := p.SSHCommand(fmt.Sprintf(certXferCmd, string(certContents), path.Join(targetPath, certFile))); err != nil {
		return err
	}

	return nil
}

func fixPermissions(p Provisioner, certPath string, targetPath string) error {
	_, certFile := path.Split(certPath)

    /*
     * TODO: Until we start dynamically generating the configuration file,
     * these must have a known naming convention on the machine.
     */
     certFile = strings.Split(certFile, "_")[1]

	_, err := p.SSHCommand(fmt.Sprintf("sudo chmod 0400 %s", path.Join(targetPath, certFile)))
	if  err != nil {
		return err
	}

	return nil
}

func GenerateCertificates(p Provisioner, k8sOptions kubernetes.KubernetesOptions, authOptions auth.AuthOptions) (error) {
  /* Generate and install certificates. Then kick off kubernetes */
  driver := p.GetDriver()
  machine := driver.GetMachineName()
  bits := 2048  // Based on the initial configuration
  targetDir := k8sOptions.K8SCertPath
  ip, err := driver.GetIP()
  if err != nil {
    return fmt.Errorf("Error retrieving address: %s", err)
  }

  err = cert.GenerateCert(
    []string{ip, "localhost"},
    k8sOptions.K8SAPICert,
    k8sOptions.K8SAPIKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "api"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating API cert: %s", err)
  }

  err = cert.GenerateCert(
    []string{""},
    k8sOptions.K8SAdminCert,
    k8sOptions.K8SAdminKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "admin"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating Admin cert: %s", err)
  }

  err = cert.GenerateCert(
    []string{},
    k8sOptions.K8SProxyCert,
    k8sOptions.K8SProxyKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "proxy"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating proxy cert: %s", err)
  }

  /* Copy certs into place */
  log.Info("Copying certs to the remote system...")

  /* Kick off the kubernetes run */
  if _, err := p.SSHCommand(fmt.Sprintf("mkdir -p %s", targetDir)); err != nil {
    return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%q,%s,%d' |sudo tee %s", k8sOptions.K8SToken, "kuser",0,path.Join(targetDir, "tokenfile.txt"))); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAPIKey, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SAPIKey, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAPICert, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SProxyCert, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SProxyKey, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SProxyKey, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAdminCert, targetDir + "/kubelet"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAdminKey, targetDir + "/kubelet"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SAdminKey, targetDir + "/kubelet"); err != nil {
    return err
  }

  /* Copy the CA cert to a known location */
  caCertContents, err := ioutil.ReadFile(authOptions.CaCertPath)
  if err != nil {
    return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s/ca.pem", caCertContents, targetDir)); err != nil {
    return err
  }

  return nil
}

func configureKubernetes(p Provisioner, k8sOptions *kubernetes.KubernetesOptions, authOptions auth.AuthOptions) (error) {
    log.Info("Configuring kubernetes...")

  if _, err := p.SSHCommand("sudo /bin/sh /usr/local/etc/init.d/kubelet stop"); err != nil {
      log.Info("Errored while attempting to stop the kubelet: %s", err)
  }

  if err := GenerateCertificates(p, *k8sOptions, authOptions); err != nil {
    return err
  }

  /* Generate and install certificates. Then kick off kubernetes */
	driver := p.GetDriver()
	machine := driver.GetMachineName()
  targetDir := k8sOptions.K8SCertPath

  /* Generate and copy a new YAML file to the target */
  configFile, err := Generatek8sManifest(machine, targetDir)
  if err != nil {
      return err
  }

  kubeletConfig, err := GenerateKubeletConfig(machine, targetDir)
  if err != nil {
      return err
  }

  /* TOOD: The target manifest directory should be a parameter throughout here */
  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s", kubeletConfig, "/etc/kubernetes/kubelet.kubeconfig")); err != nil {
      return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s", configFile, "/etc/kubernetes/manifests/kubernetes.yaml")); err != nil {
      return err
  }

	/* Lastly, start the kubelet */
    if _, err := p.SSHCommand("sudo /bin/sh /usr/local/etc/init.d/kubelet start"); err != nil {
        return err
    }

    return nil
}

func GenerateKubeletConfig(name string, targetDir string) (string, error) {
    type ConfigDetails struct {
        ClusterName   string
        CertDir       string
    }

    details := ConfigDetails{name, targetDir}
    var result bytes.Buffer

    kubeletConfigTmpl := `apiVersion: v1
kind: Config
clusters:
  - cluster:
      certificate-authority: {{.CertDir}}/ca.pem
      server: https://127.0.0.1:6443
    name: {{.ClusterName}}
contexts:
  - context:
      cluster: {{.ClusterName}}
      user: kubelet
    name: {{.ClusterName}}
users:
  - name: kubelet
    user:
      client-certificate: {{.CertDir}}/kubelet/cert.pem
      client-key: {{.CertDir}}/kubelet/key.pem`

    t, err := template.New("kubeletConfigTmpl").Parse(kubeletConfigTmpl)
    if err != nil {
        return "", err
    }

    err = t.Execute(&result, details)

    return result.String(), err

}

func Generatek8sManifest(name string, targetDir string) (string, error) {
    type ConfigDetails struct {
        ClusterName   string
        CertDir       string
    }

    details := ConfigDetails{name, targetDir}
    var result bytes.Buffer

    k8sConfigTmpl := `apiVersion: v1
kind: Pod
clusters:
  - cluster:
      certificate-authority: {{.CertDir}}/ca.pem
metadata:
  name: {{.ClusterName}}
spec:
  hostNetwork: true
  volumes:
    - name: "certs"
      hostPath:
        path: "{{.CertDir}}"
    - name: "policies"
      hostPath:
        path: "/etc/kubernetes/policies"
  containers:
    - name: "etcd"
      image: "b.gcr.io/kuar/etcd:2.1.1"
      args:
        - "--data-dir=/var/lib/etcd"
        - "--advertise-client-urls=http://127.0.0.1:2379"
        - "--listen-client-urls=http://127.0.0.1:2379"
        - "--listen-peer-urls=http://127.0.0.1:2380"
        - "--name=etcd"
    - name: "controller-manager"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      args:
        - "/hyperkube"
        - "controller-manager"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"
    - name: "apiserver"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      volumeMounts:
        - name: "certs"
          mountPath: "{{.CertDir}}"
          readOnly: true
        - name: "policies"
          mountPath: "/etc/kubernetes/policies"
          readOnly: true
      args:
        - "/hyperkube"
        - "apiserver"
        - "--authorization-mode=AlwaysAllow"
        - "--client-ca-file=/var/run/kubernetes/ca.pem"
        - "--token-auth-file={{.CertDir}}/tokenfile.txt"
        - "--allow-privileged=true"
        - "--service-cluster-ip-range=10.0.20.0/24"
        - "--insecure-bind-address=0.0.0.0"
        - "--insecure-port=8080"
        - "--secure-port=6443"
        - "--etcd-servers=http://127.0.0.1:2379"
        - "--tls-cert-file={{.CertDir}}/apiserver/cert.pem"
        - "--tls-private-key-file={{.CertDir}}/apiserver/key.pem"
        - "--v=2"
    - name: "proxy"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      securityContext:
        privileged: true
      args:
        - "/hyperkube"
        - "proxy"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"
    - name: "scheduler"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      args:
        - "/hyperkube"
        - "scheduler"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"

`
    t, err := template.New("k8sConfig").Parse(k8sConfigTmpl)
    if err != nil {
        return "", err
    }

    err = t.Execute(&result, details)

    return result.String(), err
}
<commit_msg>Forced the creation of the /etc/kubernetes directory if it doesn't exist<commit_after>package provision

import (
    "bytes"
    "fmt"
	"io/ioutil"
	"path"
    "strings"
    "text/template"

    "github.com/docker/machine/libmachine/auth"
    "github.com/docker/machine/libmachine/cert"
    "github.com/docker/machine/libmachine/kubernetes"
    "github.com/docker/machine/libmachine/log"
)

func xferCert(p Provisioner, certPath string, targetPath string) error {
	certXferCmd := "printf '%%s' '%s' | sudo tee %s"

	certContents, err := ioutil.ReadFile(certPath)
	if err != nil {
		return err
	}

    if _, err := p.SSHCommand(fmt.Sprintf("mkdir -p %s", targetPath)); err != nil {
      return err
    }

    /*
     * TODO: Until we start dynamically generating the configuration file,
     * these must have a known naming convention on the machine.
     */
     _, certFile := path.Split(certPath)
     certFile = strings.Split(certFile, "_")[1]

	if _, err := p.SSHCommand(fmt.Sprintf(certXferCmd, string(certContents), path.Join(targetPath, certFile))); err != nil {
		return err
	}

	return nil
}

func fixPermissions(p Provisioner, certPath string, targetPath string) error {
	_, certFile := path.Split(certPath)

    /*
     * TODO: Until we start dynamically generating the configuration file,
     * these must have a known naming convention on the machine.
     */
     certFile = strings.Split(certFile, "_")[1]

	_, err := p.SSHCommand(fmt.Sprintf("sudo chmod 0400 %s", path.Join(targetPath, certFile)))
	if  err != nil {
		return err
	}

	return nil
}

func GenerateCertificates(p Provisioner, k8sOptions kubernetes.KubernetesOptions, authOptions auth.AuthOptions) (error) {
  /* Generate and install certificates. Then kick off kubernetes */
  driver := p.GetDriver()
  machine := driver.GetMachineName()
  bits := 2048  // Based on the initial configuration
  targetDir := k8sOptions.K8SCertPath
  ip, err := driver.GetIP()
  if err != nil {
    return fmt.Errorf("Error retrieving address: %s", err)
  }

  err = cert.GenerateCert(
    []string{ip, "localhost"},
    k8sOptions.K8SAPICert,
    k8sOptions.K8SAPIKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "api"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating API cert: %s", err)
  }

  err = cert.GenerateCert(
    []string{""},
    k8sOptions.K8SAdminCert,
    k8sOptions.K8SAdminKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "admin"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating Admin cert: %s", err)
  }

  err = cert.GenerateCert(
    []string{},
    k8sOptions.K8SProxyCert,
    k8sOptions.K8SProxyKey,
    authOptions.CaCertPath,
    authOptions.CaPrivateKeyPath,
    kubernetes.GenOrg(machine, "proxy"),
    bits)

  if err != nil {
    return fmt.Errorf("Error generating proxy cert: %s", err)
  }

  /* Copy certs into place */
  log.Info("Copying certs to the remote system...")

  /* Kick off the kubernetes run */
  if _, err := p.SSHCommand(fmt.Sprintf("mkdir -p %s", targetDir)); err != nil {
    return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%q,%s,%d' |sudo tee %s", k8sOptions.K8SToken, "kuser",0,path.Join(targetDir, "tokenfile.txt"))); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAPIKey, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SAPIKey, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAPICert, targetDir + "/apiserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SProxyCert, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SProxyKey, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SProxyKey, targetDir + "/proxyserver"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAdminCert, targetDir + "/kubelet"); err != nil {
    return err
  }

  if err := xferCert(p, k8sOptions.K8SAdminKey, targetDir + "/kubelet"); err != nil {
    return err
  }

  if err := fixPermissions(p, k8sOptions.K8SAdminKey, targetDir + "/kubelet"); err != nil {
    return err
  }

  /* Copy the CA cert to a known location */
  caCertContents, err := ioutil.ReadFile(authOptions.CaCertPath)
  if err != nil {
    return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s/ca.pem", caCertContents, targetDir)); err != nil {
    return err
  }

  return nil
}

func configureKubernetes(p Provisioner, k8sOptions *kubernetes.KubernetesOptions, authOptions auth.AuthOptions) (error) {
    log.Info("Configuring kubernetes...")

  if _, err := p.SSHCommand("sudo /bin/sh /usr/local/etc/init.d/kubelet stop"); err != nil {
      log.Info("Errored while attempting to stop the kubelet: %s", err)
  }

  if err := GenerateCertificates(p, *k8sOptions, authOptions); err != nil {
    return err
  }

  /* Generate and install certificates. Then kick off kubernetes */
	driver := p.GetDriver()
	machine := driver.GetMachineName()
  targetDir := k8sOptions.K8SCertPath

  /* Generate and copy a new YAML file to the target */
  configFile, err := Generatek8sManifest(machine, targetDir)
  if err != nil {
      return err
  }

  kubeletConfig, err := GenerateKubeletConfig(machine, targetDir)
  if err != nil {
      return err
  }

  /* TOOD: The target manifest directory should be a parameter throughout here */
  /* Ensure that the kubernetes configuration directory exists */
  if _, err := p.SSHCommand(fmt.Sprintf("sudo mkdir -p /etc/kubernetes/manifests")); err != nil {
      return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s", kubeletConfig, "/etc/kubernetes/kubelet.kubeconfig")); err != nil {
      return err
  }

  if _, err := p.SSHCommand(fmt.Sprintf("printf '%%s' '%s' | sudo tee %s", configFile, "/etc/kubernetes/manifests/kubernetes.yaml")); err != nil {
      return err
  }

	/* Lastly, start the kubelet */
    if _, err := p.SSHCommand("sudo /bin/sh /usr/local/etc/init.d/kubelet start"); err != nil {
        return err
    }

    return nil
}

func GenerateKubeletConfig(name string, targetDir string) (string, error) {
    type ConfigDetails struct {
        ClusterName   string
        CertDir       string
    }

    details := ConfigDetails{name, targetDir}
    var result bytes.Buffer

    kubeletConfigTmpl := `apiVersion: v1
kind: Config
clusters:
  - cluster:
      certificate-authority: {{.CertDir}}/ca.pem
      server: https://127.0.0.1:6443
    name: {{.ClusterName}}
contexts:
  - context:
      cluster: {{.ClusterName}}
      user: kubelet
    name: {{.ClusterName}}
users:
  - name: kubelet
    user:
      client-certificate: {{.CertDir}}/kubelet/cert.pem
      client-key: {{.CertDir}}/kubelet/key.pem`

    t, err := template.New("kubeletConfigTmpl").Parse(kubeletConfigTmpl)
    if err != nil {
        return "", err
    }

    err = t.Execute(&result, details)

    return result.String(), err

}

func Generatek8sManifest(name string, targetDir string) (string, error) {
    type ConfigDetails struct {
        ClusterName   string
        CertDir       string
    }

    details := ConfigDetails{name, targetDir}
    var result bytes.Buffer

    k8sConfigTmpl := `apiVersion: v1
kind: Pod
clusters:
  - cluster:
      certificate-authority: {{.CertDir}}/ca.pem
metadata:
  name: {{.ClusterName}}
spec:
  hostNetwork: true
  volumes:
    - name: "certs"
      hostPath:
        path: "{{.CertDir}}"
    - name: "policies"
      hostPath:
        path: "/etc/kubernetes/policies"
  containers:
    - name: "etcd"
      image: "b.gcr.io/kuar/etcd:2.1.1"
      args:
        - "--data-dir=/var/lib/etcd"
        - "--advertise-client-urls=http://127.0.0.1:2379"
        - "--listen-client-urls=http://127.0.0.1:2379"
        - "--listen-peer-urls=http://127.0.0.1:2380"
        - "--name=etcd"
    - name: "controller-manager"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      args:
        - "/hyperkube"
        - "controller-manager"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"
    - name: "apiserver"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      volumeMounts:
        - name: "certs"
          mountPath: "{{.CertDir}}"
          readOnly: true
        - name: "policies"
          mountPath: "/etc/kubernetes/policies"
          readOnly: true
      args:
        - "/hyperkube"
        - "apiserver"
        - "--authorization-mode=AlwaysAllow"
        - "--client-ca-file=/var/run/kubernetes/ca.pem"
        - "--token-auth-file={{.CertDir}}/tokenfile.txt"
        - "--allow-privileged=true"
        - "--service-cluster-ip-range=10.0.20.0/24"
        - "--insecure-bind-address=0.0.0.0"
        - "--insecure-port=8080"
        - "--secure-port=6443"
        - "--etcd-servers=http://127.0.0.1:2379"
        - "--tls-cert-file={{.CertDir}}/apiserver/cert.pem"
        - "--tls-private-key-file={{.CertDir}}/apiserver/key.pem"
        - "--v=2"
    - name: "proxy"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      securityContext:
        privileged: true
      args:
        - "/hyperkube"
        - "proxy"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"
    - name: "scheduler"
      image: "gcr.io/google_containers/hyperkube:v1.0.3"
      args:
        - "/hyperkube"
        - "scheduler"
        - "--master=http://127.0.0.1:8080"
        - "--v=2"

`
    t, err := template.New("k8sConfig").Parse(k8sConfigTmpl)
    if err != nil {
        return "", err
    }

    err = t.Execute(&result, details)

    return result.String(), err
}
<|endoftext|>


<commit_before>package lints

import (
	"regexp"

	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint/util"
)

type DNSNameProperCharacters struct {
	CompiledExpression *regexp.Regexp
}

func (l *DNSNameProperCharacters) Initialize() error {
	const dnsNameRegexp = `^(\*\.)?(\?\.)*(A-Za-z0-9*_-]+\.)*[A-Za-z0-9*_-]*$`
	var err error
	l.CompiledExpression, err = regexp.Compile(dnsNameRegexp)

	return err
}

func (l *DNSNameProperCharacters) CheckApplies(c *x509.Certificate) bool {
	return util.IsSubscriberCert(c) && util.DNSNamesExist(c)
}

func (l *DNSNameProperCharacters) Execute(c *x509.Certificate) *LintResult {
	if c.Subject.CommonName != "" {
		if !l.CompiledExpression.MatchString(c.Subject.CommonName) {
			return &LintResult{Status: Error}
		}
	}
	for _, dns := range c.DNSNames {
		if !l.CompiledExpression.MatchString(dns) {
			return &LintResult{Status: Error}
		}
	}
	return &LintResult{Status: Pass}
}

func init() {
	RegisterLint(&Lint{
		Name:          "e_dnsname_bad_character_in_label",
		Description:   "Characters in labels of DNSNames MUST be alphanumeric, - , _ or *",
		Citation:      "BRs: 7.1.4.2",
		Source:        CABFBaselineRequirements,
		EffectiveDate: util.CABEffectiveDate,
		Lint:          &DNSNameProperCharacters{},
	})
}
<commit_msg>Add missing [ to dnsNameRegexp. (#181)<commit_after>package lints

import (
	"regexp"

	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint/util"
)

type DNSNameProperCharacters struct {
	CompiledExpression *regexp.Regexp
}

func (l *DNSNameProperCharacters) Initialize() error {
	const dnsNameRegexp = `^(\*\.)?(\?\.)*([A-Za-z0-9*_-]+\.)*[A-Za-z0-9*_-]*$`
	var err error
	l.CompiledExpression, err = regexp.Compile(dnsNameRegexp)

	return err
}

func (l *DNSNameProperCharacters) CheckApplies(c *x509.Certificate) bool {
	return util.IsSubscriberCert(c) && util.DNSNamesExist(c)
}

func (l *DNSNameProperCharacters) Execute(c *x509.Certificate) *LintResult {
	if c.Subject.CommonName != "" {
		if !l.CompiledExpression.MatchString(c.Subject.CommonName) {
			return &LintResult{Status: Error}
		}
	}
	for _, dns := range c.DNSNames {
		if !l.CompiledExpression.MatchString(dns) {
			return &LintResult{Status: Error}
		}
	}
	return &LintResult{Status: Pass}
}

func init() {
	RegisterLint(&Lint{
		Name:          "e_dnsname_bad_character_in_label",
		Description:   "Characters in labels of DNSNames MUST be alphanumeric, - , _ or *",
		Citation:      "BRs: 7.1.4.2",
		Source:        CABFBaselineRequirements,
		EffectiveDate: util.CABEffectiveDate,
		Lint:          &DNSNameProperCharacters{},
	})
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"
)

const listenProxyAddr = "127.0.0.1:9997"

var globalClient *Client

var testURLs = map[string][]byte{
	"http://www.google.com/humans.txt": []byte("Google is built by a large team of engineers, designers, researchers, robots, and others in many different sites across the globe. It is updated continuously, and built with more tools and technologies than we can shake a stick at. If you'd like to help us out, see google.com/careers.\n"),
	// TODO: This is not working, we actually need to implement a CONNECT request to proxy HTTPs traffic.
	// "https://www.google.com/humans.txt": []byte("Google is built by a large team of engineers, designers, researchers, robots, and others in many different sites across the globe. It is updated continuously, and built with more tools and technologies than we can shake a stick at. If you'd like to help us out, see google.com/careers.\n"),
}

// Attempt to create a server in a goroutine and stop it from other place.
func TestListenAndServeStop(t *testing.T) {

	// Creating a client.
	c := NewClient(listenProxyAddr)

	if c == nil {
		t.Fatal("You should be able to create a client.")
	}

	// Make the client listen on a goroutine.
	go func() {
		c.ListenAndServe()
	}()

	// Allow it some seconds to start.
	time.Sleep(time.Millisecond * 100)

	// Attempt to stop server.
	if err := c.Stop(); err != nil {
		t.Fatal("You should be able to close listening client.")
	}

}

func TestListenAndServeAgain(t *testing.T) {
	// Since we've closed out server, we should be able to launch another at the
	// same address.

	go func() {
		globalClient = NewClient(listenProxyAddr)

		if err := globalClient.ListenAndServe(); err != nil {
			t.Fatal(err)
		}

	}()

	// Allow it some seconds to start.
	time.Sleep(time.Millisecond * 100)
}

func TestListenAndServeProxy(t *testing.T) {
	var wg sync.WaitGroup

	// Testing the client we've just opened.
	for uri, expectedContent := range testURLs {
		wg.Add(1)

		go func(wg *sync.WaitGroup) {
			if err := testReverseProxy(uri, expectedContent); err != nil {
				t.Fatal(err)
			}
			wg.Done()
		}(&wg)

	}

	wg.Wait()
}

func TestCloseClient(t *testing.T) {

	// Closing the client that is still opened.
	if err := globalClient.Stop(); err != nil {
		t.Fatal("You should be able to close listening client.")
	}
}

func testReverseProxy(uri string, expectedContent []byte) (err error) {
	var req *http.Request
	var u *url.URL

	if u, err = url.Parse(uri); err != nil {
		return err
	}

	port := 80

	if u.Scheme == "https" {
		// TODO: implement a CONNECT request.
	}

	req = &http.Request{
		Method: "GET",
		URL: &url.URL{
			Scheme: u.Scheme,
			Host:   u.Host,
			Path:   uri,
		},
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header: http.Header{
			"Host": {fmt.Sprintf("%s:%d", u.Host, port)},
		},
	}

	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(n, a string) (net.Conn, error) {
				return net.Dial("tcp", listenProxyAddr)
			},
		},
	}

	var res *http.Response

	if res, err = client.Do(req); err != nil {
		return err
	}

	var buf []byte

	buf, err = ioutil.ReadAll(res.Body)

	fmt.Printf(string(buf))

	if bytes.Equal(buf, expectedContent) == false {
		return fmt.Errorf("The response we've got from %s differs from what we expected.", uri)
	}

	return nil
}
<commit_msg>Implementing test as in testRequest() from flashlight_test.go.<commit_after>package client

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"
)

const listenProxyAddr = "127.0.0.1:9997"

var globalClient *Client

var testURLs = map[string][]byte{
	"http://www.google.com/humans.txt":  []byte("Google is built by a large team of engineers, designers, researchers, robots, and others in many different sites across the globe. It is updated continuously, and built with more tools and technologies than we can shake a stick at. If you'd like to help us out, see google.com/careers.\n"),
	"https://www.google.com/humans.txt": []byte("Google is built by a large team of engineers, designers, researchers, robots, and others in many different sites across the globe. It is updated continuously, and built with more tools and technologies than we can shake a stick at. If you'd like to help us out, see google.com/careers.\n"),
}

// Attempt to create a server in a goroutine and stop it from other place.
func TestListenAndServeStop(t *testing.T) {

	// Creating a client.
	c := NewClient(listenProxyAddr)

	if c == nil {
		t.Fatal("You should be able to create a client.")
	}

	// Make the client listen on a goroutine.
	go func() {
		c.ListenAndServe()
	}()

	// Allow it some seconds to start.
	time.Sleep(time.Millisecond * 100)

	// Attempt to stop server.
	if err := c.Stop(); err != nil {
		t.Fatal("You should be able to close listening client.")
	}

}

func TestListenAndServeAgain(t *testing.T) {
	// Since we've closed out server, we should be able to launch another at the
	// same address.

	go func() {
		globalClient = NewClient(listenProxyAddr)

		if err := globalClient.ListenAndServe(); err != nil {
			t.Fatal(err)
		}

	}()

	// Allow it some seconds to start.
	time.Sleep(time.Millisecond * 100)
}

func TestListenAndServeProxy(t *testing.T) {
	var wg sync.WaitGroup

	// Testing the client we've just opened.
	for uri, expectedContent := range testURLs {
		wg.Add(1)

		go func(wg *sync.WaitGroup) {
			if err := testReverseProxy(uri, expectedContent); err != nil {
				t.Fatal(err)
			}
			wg.Done()
		}(&wg)

	}

	wg.Wait()
}

func TestCloseClient(t *testing.T) {

	// Closing the client that is still opened.
	if err := globalClient.Stop(); err != nil {
		t.Fatal("You should be able to close listening client.")
	}
}

func testReverseProxy(destURL string, expectedContent []byte) (err error) {
	var req *http.Request

	if req, err = http.NewRequest("GET", destURL, nil); err != nil {
		return err
	}

	client := &http.Client{
		Transport: &http.Transport{
			Proxy: func(req *http.Request) (*url.URL, error) {
				return url.Parse(listenProxyAddr)
			},
			Dial: func(n, a string) (net.Conn, error) {
				return net.Dial("tcp", listenProxyAddr)
			},
		},
	}

	var res *http.Response

	if res, err = client.Do(req); err != nil {
		return err
	}

	var buf []byte

	buf, err = ioutil.ReadAll(res.Body)

	fmt.Printf(string(buf))

	if bytes.Equal(buf, expectedContent) == false {
		return fmt.Errorf("The response we've got from %s differs from what we expected.", destURL)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"go/parser"
	"go/ast"
	"go/token"
)

type AutoCompleteFile struct {
	name        string
	packageName string

	decls     map[string]*Decl
	packages  PackageImports
	filescope *Scope
	scope     *Scope

	cursor int // for current file buffer only
}

func NewPackageFile(name string) *AutoCompleteFile {
	p := new(AutoCompleteFile)
	p.name = name
	p.cursor = -1
	return p
}

// this one is used for current file buffer exclusively
func (f *AutoCompleteFile) processData(data []byte) {
	cur, filedata, block := RipOffDecl(data, f.cursor)
	file, _ := parser.ParseFile("", filedata, 0)
	f.packageName = packageName(file)

	f.decls = make(map[string]*Decl)
	f.packages = NewPackageImports(f.name, file.Decls)
	f.filescope = NewScope(nil)
	f.scope = f.filescope

	anonymifyAst(file.Decls, 0, f.filescope)

	// process all top-level declarations
	for _, decl := range file.Decls {
		appendToTopDecls(f.decls, decl, f.scope)
	}
	if block != nil {
		// process local function as top-level declaration
		decls, _ := parser.ParseDeclList("", block)

		anonymifyAst(decls, 0, f.filescope)

		for _, decl := range decls {
			appendToTopDecls(f.decls, decl, f.scope)
		}

		// process function internals
		f.cursor = cur
		for _, decl := range decls {
			f.processDeclLocals(decl)
		}
	}

}

func (f *AutoCompleteFile) processDeclLocals(decl ast.Decl) {
	switch t := decl.(type) {
	case *ast.FuncDecl:
		if f.cursorIn(t.Body) {
			s := f.scope
			f.scope = NewScope(f.scope)

			f.processFieldList(t.Recv, s)
			f.processFieldList(t.Type.Params, s)
			f.processFieldList(t.Type.Results, s)
			f.processBlockStmt(t.Body)

		}
	}
}

func (f *AutoCompleteFile) processDecl(decl ast.Decl) {
	if t, ok := decl.(*ast.GenDecl); ok && t.Offset > f.cursor {
		return
	}
	foreachDecl(decl, func(data *foreachDeclStruct) {
		class := astDeclClass(data.decl)
		for i, name := range data.names {
			typ, v, vi := data.typeValueIndex(i, 0)

			d := NewDecl2(name.Name, class, 0, typ, v, vi, f.scope)
			if d == nil {
				return
			}

			// the declaration itself has a scope which follows it's definition
			// and it's false for type declarations
			// TODO: WTF?
			if d.Class != DECL_TYPE {
				f.scope = NewScope(f.scope)
			}
			f.scope.addNamedDecl(d)
		}
	})
}

func (f *AutoCompleteFile) processBlockStmt(block *ast.BlockStmt) {
	if block != nil && f.cursorIn(block) {
		f.scope = AdvanceScope(f.scope)

		for _, stmt := range block.List {
			f.processStmt(stmt)
		}

		// hack to process all func literals
		v := new(funcLitVisitor)
		v.ctx = f
		ast.Walk(v, block)
	}
}

type funcLitVisitor struct {
	ctx *AutoCompleteFile
}

func (v *funcLitVisitor) Visit(node interface{}) ast.Visitor {
	if t, ok := node.(*ast.FuncLit); ok && v.ctx.cursorIn(t.Body) {
		s := v.ctx.scope
		v.ctx.scope = AdvanceScope(v.ctx.scope)

		v.ctx.processFieldList(t.Type.Params, s)
		v.ctx.processFieldList(t.Type.Results, s)
		v.ctx.processBlockStmt(t.Body)

		return nil
	}
	return v
}

func (f *AutoCompleteFile) processStmt(stmt ast.Stmt) {
	switch t := stmt.(type) {
	case *ast.DeclStmt:
		f.processDecl(t.Decl)
	case *ast.AssignStmt:
		f.processAssignStmt(t)
	case *ast.IfStmt:
		if f.cursorIn(t.Body) {
			f.scope = AdvanceScope(f.scope)

			f.processStmt(t.Init)
			f.processBlockStmt(t.Body)
		}
		f.processStmt(t.Else)
	case *ast.BlockStmt:
		f.processBlockStmt(t)
	case *ast.RangeStmt:
		f.processRangeStmt(t)
	case *ast.ForStmt:
		if f.cursorIn(t.Body) {
			f.scope = AdvanceScope(f.scope)

			f.processStmt(t.Init)
			f.processBlockStmt(t.Body)
		}
	case *ast.SwitchStmt:
		f.processSwitchStmt(t)
	case *ast.TypeSwitchStmt:
		f.processTypeSwitchStmt(t)
	case *ast.SelectStmt:
		f.processSelectStmt(t)
	case *ast.LabeledStmt:
		f.processStmt(t.Stmt)
	}
}

func (f *AutoCompleteFile) processSelectStmt(a *ast.SelectStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	var lastCursorAfter *ast.CommClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.CommClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}

	if lastCursorAfter != nil {
		if lastCursorAfter.Lhs != nil && lastCursorAfter.Tok == token.DEFINE {
			vname := lastCursorAfter.Lhs.(*ast.Ident).Name
			v := NewDeclVar(vname, nil, lastCursorAfter.Rhs, -1, f.scope)
			f.scope.addNamedDecl(v)
		}
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processTypeSwitchStmt(a *ast.TypeSwitchStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	f.processStmt(a.Init)
	// type var
	var tv *Decl
	if a, ok := a.Assign.(*ast.AssignStmt); ok {
		lhs := a.Lhs
		rhs := a.Rhs
		if lhs != nil && len(lhs) == 1 {
			tvname := lhs[0].(*ast.Ident).Name
			tv = NewDeclVar(tvname, nil, rhs[0], -1, f.scope)
		}
	}

	var lastCursorAfter *ast.TypeCaseClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.TypeCaseClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}

	if lastCursorAfter != nil {
		if tv != nil {
			if lastCursorAfter.Types != nil && len(lastCursorAfter.Types) == 1 {
				tv.Type = lastCursorAfter.Types[0]
				tv.Value = nil
			}
			f.scope.addNamedDecl(tv)
		}
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processSwitchStmt(a *ast.SwitchStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	f.processStmt(a.Init)
	var lastCursorAfter *ast.CaseClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.CaseClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}
	if lastCursorAfter != nil {
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processRangeStmt(a *ast.RangeStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	if a.Tok == token.DEFINE {
		if t, ok := a.Key.(*ast.Ident); ok {
			d := NewDeclVar(t.Name, nil, a.X, 0, f.scope)
			if d != nil {
				d.Flags |= DECL_RANGEVAR
				f.scope.addNamedDecl(d)
			}
		}

		if a.Value != nil {
			if t, ok := a.Value.(*ast.Ident); ok {
				d := NewDeclVar(t.Name, nil, a.X, 1, f.scope)
				if d != nil {
					d.Flags |= DECL_RANGEVAR
					f.scope.addNamedDecl(d)
				}
			}
		}
	}

	f.processBlockStmt(a.Body)
}

func (f *AutoCompleteFile) processAssignStmt(a *ast.AssignStmt) {
	if a.Tok != token.DEFINE || a.TokPos.Offset > f.cursor {
		return
	}

	names := make([]*ast.Ident, len(a.Lhs))
	for i, name := range a.Lhs {
		id, ok := name.(*ast.Ident)
		if !ok {
			// something is wrong, just ignore the whole stmt
			return
		}
		names[i] = id
	}

	pack := declPack{names, nil, a.Rhs}
	for i, name := range pack.names {
		typ, v, vi := pack.typeValueIndex(i, 0)
		d := NewDeclVar(name.Name, typ, v, vi, f.scope)
		if d == nil {
			continue
		}

		f.scope = NewScope(f.scope)
		f.scope.addNamedDecl(d)
	}
}

func (f *AutoCompleteFile) processFieldList(fieldList *ast.FieldList, s *Scope) {
	if fieldList != nil {
		decls := astFieldListToDecls(fieldList, DECL_VAR, 0, s)
		for _, d := range decls {
			f.scope.addNamedDecl(d)
		}
	}
}

func (f *AutoCompleteFile) cursorIn(block *ast.BlockStmt) bool {
	if f.cursor == -1 || block == nil {
		return false
	}

	if f.cursor >= block.Offset && f.cursor <= block.Rbrace.Offset {
		return true
	}
	return false
}
<commit_msg>Fix TODO item.<commit_after>package main

import (
	"go/parser"
	"go/ast"
	"go/token"
)

type AutoCompleteFile struct {
	name        string
	packageName string

	decls     map[string]*Decl
	packages  PackageImports
	filescope *Scope
	scope     *Scope

	cursor int // for current file buffer only
}

func NewPackageFile(name string) *AutoCompleteFile {
	p := new(AutoCompleteFile)
	p.name = name
	p.cursor = -1
	return p
}

// this one is used for current file buffer exclusively
func (f *AutoCompleteFile) processData(data []byte) {
	cur, filedata, block := RipOffDecl(data, f.cursor)
	file, _ := parser.ParseFile("", filedata, 0)
	f.packageName = packageName(file)

	f.decls = make(map[string]*Decl)
	f.packages = NewPackageImports(f.name, file.Decls)
	f.filescope = NewScope(nil)
	f.scope = f.filescope

	anonymifyAst(file.Decls, 0, f.filescope)

	// process all top-level declarations
	for _, decl := range file.Decls {
		appendToTopDecls(f.decls, decl, f.scope)
	}
	if block != nil {
		// process local function as top-level declaration
		decls, _ := parser.ParseDeclList("", block)

		anonymifyAst(decls, 0, f.filescope)

		for _, decl := range decls {
			appendToTopDecls(f.decls, decl, f.scope)
		}

		// process function internals
		f.cursor = cur
		for _, decl := range decls {
			f.processDeclLocals(decl)
		}
	}

}

func (f *AutoCompleteFile) processDeclLocals(decl ast.Decl) {
	switch t := decl.(type) {
	case *ast.FuncDecl:
		if f.cursorIn(t.Body) {
			s := f.scope
			f.scope = NewScope(f.scope)

			f.processFieldList(t.Recv, s)
			f.processFieldList(t.Type.Params, s)
			f.processFieldList(t.Type.Results, s)
			f.processBlockStmt(t.Body)

		}
	}
}

func (f *AutoCompleteFile) processDecl(decl ast.Decl) {
	if t, ok := decl.(*ast.GenDecl); ok && t.Offset > f.cursor {
		return
	}
	foreachDecl(decl, func(data *foreachDeclStruct) {
		class := astDeclClass(data.decl)
		for i, name := range data.names {
			typ, v, vi := data.typeValueIndex(i, 0)

			d := NewDecl2(name.Name, class, 0, typ, v, vi, f.scope)
			if d == nil {
				return
			}

			// TODO: The code here is incorrect, but I wasn't able to break
			// it. Maybe it is correct for autocompletion case.
			if d.Class != DECL_TYPE {
				f.scope = NewScope(f.scope)
			}
			f.scope.addNamedDecl(d)
		}
	})
}

func (f *AutoCompleteFile) processBlockStmt(block *ast.BlockStmt) {
	if block != nil && f.cursorIn(block) {
		f.scope = AdvanceScope(f.scope)

		for _, stmt := range block.List {
			f.processStmt(stmt)
		}

		// hack to process all func literals
		v := new(funcLitVisitor)
		v.ctx = f
		ast.Walk(v, block)
	}
}

type funcLitVisitor struct {
	ctx *AutoCompleteFile
}

func (v *funcLitVisitor) Visit(node interface{}) ast.Visitor {
	if t, ok := node.(*ast.FuncLit); ok && v.ctx.cursorIn(t.Body) {
		s := v.ctx.scope
		v.ctx.scope = AdvanceScope(v.ctx.scope)

		v.ctx.processFieldList(t.Type.Params, s)
		v.ctx.processFieldList(t.Type.Results, s)
		v.ctx.processBlockStmt(t.Body)

		return nil
	}
	return v
}

func (f *AutoCompleteFile) processStmt(stmt ast.Stmt) {
	switch t := stmt.(type) {
	case *ast.DeclStmt:
		f.processDecl(t.Decl)
	case *ast.AssignStmt:
		f.processAssignStmt(t)
	case *ast.IfStmt:
		if f.cursorIn(t.Body) {
			f.scope = AdvanceScope(f.scope)

			f.processStmt(t.Init)
			f.processBlockStmt(t.Body)
		}
		f.processStmt(t.Else)
	case *ast.BlockStmt:
		f.processBlockStmt(t)
	case *ast.RangeStmt:
		f.processRangeStmt(t)
	case *ast.ForStmt:
		if f.cursorIn(t.Body) {
			f.scope = AdvanceScope(f.scope)

			f.processStmt(t.Init)
			f.processBlockStmt(t.Body)
		}
	case *ast.SwitchStmt:
		f.processSwitchStmt(t)
	case *ast.TypeSwitchStmt:
		f.processTypeSwitchStmt(t)
	case *ast.SelectStmt:
		f.processSelectStmt(t)
	case *ast.LabeledStmt:
		f.processStmt(t.Stmt)
	}
}

func (f *AutoCompleteFile) processSelectStmt(a *ast.SelectStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	var lastCursorAfter *ast.CommClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.CommClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}

	if lastCursorAfter != nil {
		if lastCursorAfter.Lhs != nil && lastCursorAfter.Tok == token.DEFINE {
			vname := lastCursorAfter.Lhs.(*ast.Ident).Name
			v := NewDeclVar(vname, nil, lastCursorAfter.Rhs, -1, f.scope)
			f.scope.addNamedDecl(v)
		}
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processTypeSwitchStmt(a *ast.TypeSwitchStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	f.processStmt(a.Init)
	// type var
	var tv *Decl
	if a, ok := a.Assign.(*ast.AssignStmt); ok {
		lhs := a.Lhs
		rhs := a.Rhs
		if lhs != nil && len(lhs) == 1 {
			tvname := lhs[0].(*ast.Ident).Name
			tv = NewDeclVar(tvname, nil, rhs[0], -1, f.scope)
		}
	}

	var lastCursorAfter *ast.TypeCaseClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.TypeCaseClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}

	if lastCursorAfter != nil {
		if tv != nil {
			if lastCursorAfter.Types != nil && len(lastCursorAfter.Types) == 1 {
				tv.Type = lastCursorAfter.Types[0]
				tv.Value = nil
			}
			f.scope.addNamedDecl(tv)
		}
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processSwitchStmt(a *ast.SwitchStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	f.processStmt(a.Init)
	var lastCursorAfter *ast.CaseClause
	for _, s := range a.Body.List {
		if cc := s.(*ast.CaseClause); f.cursor > cc.Colon.Offset {
			lastCursorAfter = cc
		}
	}
	if lastCursorAfter != nil {
		for _, s := range lastCursorAfter.Body {
			f.processStmt(s)
		}
	}
}

func (f *AutoCompleteFile) processRangeStmt(a *ast.RangeStmt) {
	if !f.cursorIn(a.Body) {
		return
	}
	f.scope = AdvanceScope(f.scope)

	if a.Tok == token.DEFINE {
		if t, ok := a.Key.(*ast.Ident); ok {
			d := NewDeclVar(t.Name, nil, a.X, 0, f.scope)
			if d != nil {
				d.Flags |= DECL_RANGEVAR
				f.scope.addNamedDecl(d)
			}
		}

		if a.Value != nil {
			if t, ok := a.Value.(*ast.Ident); ok {
				d := NewDeclVar(t.Name, nil, a.X, 1, f.scope)
				if d != nil {
					d.Flags |= DECL_RANGEVAR
					f.scope.addNamedDecl(d)
				}
			}
		}
	}

	f.processBlockStmt(a.Body)
}

func (f *AutoCompleteFile) processAssignStmt(a *ast.AssignStmt) {
	if a.Tok != token.DEFINE || a.TokPos.Offset > f.cursor {
		return
	}

	names := make([]*ast.Ident, len(a.Lhs))
	for i, name := range a.Lhs {
		id, ok := name.(*ast.Ident)
		if !ok {
			// something is wrong, just ignore the whole stmt
			return
		}
		names[i] = id
	}

	pack := declPack{names, nil, a.Rhs}
	for i, name := range pack.names {
		typ, v, vi := pack.typeValueIndex(i, 0)
		d := NewDeclVar(name.Name, typ, v, vi, f.scope)
		if d == nil {
			continue
		}

		f.scope = NewScope(f.scope)
		f.scope.addNamedDecl(d)
	}
}

func (f *AutoCompleteFile) processFieldList(fieldList *ast.FieldList, s *Scope) {
	if fieldList != nil {
		decls := astFieldListToDecls(fieldList, DECL_VAR, 0, s)
		for _, d := range decls {
			f.scope.addNamedDecl(d)
		}
	}
}

func (f *AutoCompleteFile) cursorIn(block *ast.BlockStmt) bool {
	if f.cursor == -1 || block == nil {
		return false
	}

	if f.cursor >= block.Offset && f.cursor <= block.Rbrace.Offset {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>/*
Copyright 2021 Gravitational, Inc.

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

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

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

// Package sshutils defines several functions and types used across the
// Teleport API and other Teleport packages when working with SSH.
package sshutils

import (
	"crypto/subtle"
	"fmt"
	"net"
	"runtime"

	"github.com/gravitational/teleport/api/constants"
	"github.com/gravitational/teleport/api/defaults"

	"github.com/gravitational/trace"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

// ParseCertificate parses an SSH certificate from the authorized_keys format.
func ParseCertificate(buf []byte) (*ssh.Certificate, error) {
	k, _, _, _, err := ssh.ParseAuthorizedKey(buf)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	cert, ok := k.(*ssh.Certificate)
	if !ok {
		return nil, trace.BadParameter("not an SSH certificate")
	}

	return cert, nil
}

// ProxyClientSSHConfig returns an ssh.ClientConfig with SSH credentials from this
// Key and HostKeyCallback matching SSH CAs in the Key.
//
// The config is set up to authenticate to proxy with the first available principal.
//
func ProxyClientSSHConfig(sshCert, privKey []byte, caCerts [][]byte) (*ssh.ClientConfig, error) {
	cert, err := ParseCertificate(sshCert)
	if err != nil {
		return nil, trace.Wrap(err, "failed to extract username from SSH certificate")
	}

	authMethod, err := AsAuthMethod(cert, privKey)
	if err != nil {
		return nil, trace.Wrap(err, "failed to convert key pair to auth method")
	}

	hostKeyCallback, err := HostKeyCallback(caCerts)
	if err != nil {
		return nil, trace.Wrap(err, "failed to convert certificate authorities to HostKeyCallback")
	}

	// The KeyId is not always a valid principal, so we use the first valid principal instead.
	user := cert.KeyId
	if len(cert.ValidPrincipals) > 0 {
		user = cert.ValidPrincipals[0]
	}

	return &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{authMethod},
		HostKeyCallback: hostKeyCallback,
		Timeout:         defaults.DefaultDialTimeout,
	}, nil
}

// AsSigner returns an ssh.Signer from raw marshaled key and certificate.
func AsSigner(sshCert *ssh.Certificate, privKey []byte) (ssh.Signer, error) {
	keys, err := AsAgentKeys(sshCert, privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	signer, err := ssh.NewSignerFromKey(keys[0].PrivateKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	signer, err = ssh.NewCertSigner(keys[0].Certificate, signer)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return signer, nil
}

// AsAuthMethod returns an "auth method" interface, a common abstraction
// used by Golang SSH library. This is how you actually use a Key to feed
// it into the SSH lib.
func AsAuthMethod(sshCert *ssh.Certificate, privKey []byte) (ssh.AuthMethod, error) {
	signer, err := AsSigner(sshCert, privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return ssh.PublicKeys(signer), nil
}

// AsAgentKeys converts Key struct to a []*agent.AddedKey. All elements
// of the []*agent.AddedKey slice need to be loaded into the agent!
func AsAgentKeys(sshCert *ssh.Certificate, privKey []byte) ([]agent.AddedKey, error) {
	// unmarshal private key bytes into a *rsa.PrivateKey
	privateKey, err := ssh.ParseRawPrivateKey(privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// put a teleport identifier along with the teleport user into the comment field
	comment := fmt.Sprintf("teleport:%v", sshCert.KeyId)

	// On Windows, return the certificate with the private key embedded.
	if runtime.GOOS == constants.WindowsOS {
		return []agent.AddedKey{
			{
				PrivateKey:       privateKey,
				Certificate:      sshCert,
				Comment:          comment,
				LifetimeSecs:     0,
				ConfirmBeforeUse: false,
			},
		}, nil
	}

	// On Unix, return the certificate (with embedded private key) as well as
	// a private key.
	//
	// This is done because OpenSSH clients older than OpenSSH 7.3/7.3p1
	// (2016-08-01) have a bug in how they use certificates that have been loaded
	// in an agent. Specifically when you add a certificate to an agent, you can't
	// just embed the private key within the certificate, you have to add the
	// certificate and private key to the agent separately. Teleport works around
	// this behavior to ensure OpenSSH interoperability.
	//
	// For more details see the following: https://bugzilla.mindrot.org/show_bug.cgi?id=2550
	// WARNING: callers expect the returned slice to be __exactly as it is__
	return []agent.AddedKey{
		{
			PrivateKey:       privateKey,
			Certificate:      sshCert,
			Comment:          comment,
			LifetimeSecs:     0,
			ConfirmBeforeUse: false,
		},
		{
			PrivateKey:       privateKey,
			Certificate:      nil,
			Comment:          comment,
			LifetimeSecs:     0,
			ConfirmBeforeUse: false,
		},
	}, nil
}

// HostKeyCallback returns an ssh.HostKeyCallback that validates host
// keys/certs against SSH CAs in the Key.
//
// If not CAs are present in the Key, the returned ssh.HostKeyCallback is nil.
// This causes golang.org/x/crypto/ssh to prompt the user to verify host key
// fingerprint (same as OpenSSH does for an unknown host).
func HostKeyCallback(caCerts [][]byte) (ssh.HostKeyCallback, error) {
	var trustedKeys []ssh.PublicKey
	for _, caCert := range caCerts {
		_, _, publicKey, _, _, err := ssh.ParseKnownHosts(caCert)
		if err != nil {
			return nil, trace.BadParameter("failed parsing CA cert: %v; raw CA cert line: %q", err, caCert)
		}
		trustedKeys = append(trustedKeys, publicKey)
	}
	// No CAs are provided, return a nil callback which will prompt the user
	// for trust.
	if len(trustedKeys) == 0 {
		return nil, nil
	}

	return func(host string, a net.Addr, hostKey ssh.PublicKey) error {
		clusterCert, ok := hostKey.(*ssh.Certificate)
		if ok {
			hostKey = clusterCert.SignatureKey
		}
		for _, trustedKey := range trustedKeys {
			if KeysEqual(trustedKey, hostKey) {
				return nil
			}
		}
		return trace.AccessDenied("host %v is untrusted or Teleport CA has been rotated; try getting new credentials by logging in again ('tsh login') or re-exporting the identity file ('tctl auth sign' or 'tsh login -o'), depending on how you got them initially", host)
	}, nil
}

// KeysEqual is constant time compare of the keys to avoid timing attacks
func KeysEqual(ak, bk ssh.PublicKey) bool {
	a := ssh.Marshal(ak)
	b := ssh.Marshal(bk)
	return (len(a) == len(b) && subtle.ConstantTimeCompare(a, b) == 1)
}
<commit_msg>Fix profile credential loader known_hosts (#7532)<commit_after>/*
Copyright 2021 Gravitational, Inc.

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

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

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

// Package sshutils defines several functions and types used across the
// Teleport API and other Teleport packages when working with SSH.
package sshutils

import (
	"crypto/subtle"
	"fmt"
	"io"
	"net"
	"runtime"

	"github.com/gravitational/teleport/api/constants"
	"github.com/gravitational/teleport/api/defaults"

	"github.com/gravitational/trace"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

// ParseCertificate parses an SSH certificate from the authorized_keys format.
func ParseCertificate(buf []byte) (*ssh.Certificate, error) {
	k, _, _, _, err := ssh.ParseAuthorizedKey(buf)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	cert, ok := k.(*ssh.Certificate)
	if !ok {
		return nil, trace.BadParameter("not an SSH certificate")
	}

	return cert, nil
}

// ProxyClientSSHConfig returns an ssh.ClientConfig with SSH credentials from this
// Key and HostKeyCallback matching SSH CAs in the Key.
//
// The config is set up to authenticate to proxy with the first available principal.
//
func ProxyClientSSHConfig(sshCert, privKey []byte, caCerts [][]byte) (*ssh.ClientConfig, error) {
	cert, err := ParseCertificate(sshCert)
	if err != nil {
		return nil, trace.Wrap(err, "failed to extract username from SSH certificate")
	}

	authMethod, err := AsAuthMethod(cert, privKey)
	if err != nil {
		return nil, trace.Wrap(err, "failed to convert key pair to auth method")
	}

	hostKeyCallback, err := HostKeyCallback(caCerts)
	if err != nil {
		return nil, trace.Wrap(err, "failed to convert certificate authorities to HostKeyCallback")
	}

	// The KeyId is not always a valid principal, so we use the first valid principal instead.
	user := cert.KeyId
	if len(cert.ValidPrincipals) > 0 {
		user = cert.ValidPrincipals[0]
	}

	return &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{authMethod},
		HostKeyCallback: hostKeyCallback,
		Timeout:         defaults.DefaultDialTimeout,
	}, nil
}

// AsSigner returns an ssh.Signer from raw marshaled key and certificate.
func AsSigner(sshCert *ssh.Certificate, privKey []byte) (ssh.Signer, error) {
	keys, err := AsAgentKeys(sshCert, privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	signer, err := ssh.NewSignerFromKey(keys[0].PrivateKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	signer, err = ssh.NewCertSigner(keys[0].Certificate, signer)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return signer, nil
}

// AsAuthMethod returns an "auth method" interface, a common abstraction
// used by Golang SSH library. This is how you actually use a Key to feed
// it into the SSH lib.
func AsAuthMethod(sshCert *ssh.Certificate, privKey []byte) (ssh.AuthMethod, error) {
	signer, err := AsSigner(sshCert, privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return ssh.PublicKeys(signer), nil
}

// AsAgentKeys converts Key struct to a []*agent.AddedKey. All elements
// of the []*agent.AddedKey slice need to be loaded into the agent!
func AsAgentKeys(sshCert *ssh.Certificate, privKey []byte) ([]agent.AddedKey, error) {
	// unmarshal private key bytes into a *rsa.PrivateKey
	privateKey, err := ssh.ParseRawPrivateKey(privKey)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// put a teleport identifier along with the teleport user into the comment field
	comment := fmt.Sprintf("teleport:%v", sshCert.KeyId)

	// On Windows, return the certificate with the private key embedded.
	if runtime.GOOS == constants.WindowsOS {
		return []agent.AddedKey{
			{
				PrivateKey:       privateKey,
				Certificate:      sshCert,
				Comment:          comment,
				LifetimeSecs:     0,
				ConfirmBeforeUse: false,
			},
		}, nil
	}

	// On Unix, return the certificate (with embedded private key) as well as
	// a private key.
	//
	// This is done because OpenSSH clients older than OpenSSH 7.3/7.3p1
	// (2016-08-01) have a bug in how they use certificates that have been loaded
	// in an agent. Specifically when you add a certificate to an agent, you can't
	// just embed the private key within the certificate, you have to add the
	// certificate and private key to the agent separately. Teleport works around
	// this behavior to ensure OpenSSH interoperability.
	//
	// For more details see the following: https://bugzilla.mindrot.org/show_bug.cgi?id=2550
	// WARNING: callers expect the returned slice to be __exactly as it is__
	return []agent.AddedKey{
		{
			PrivateKey:       privateKey,
			Certificate:      sshCert,
			Comment:          comment,
			LifetimeSecs:     0,
			ConfirmBeforeUse: false,
		},
		{
			PrivateKey:       privateKey,
			Certificate:      nil,
			Comment:          comment,
			LifetimeSecs:     0,
			ConfirmBeforeUse: false,
		},
	}, nil
}

// HostKeyCallback returns an ssh.HostKeyCallback that validates host
// keys/certs against SSH CAs in the Key.
//
// If not CAs are present in the Key, the returned ssh.HostKeyCallback is nil.
// This causes golang.org/x/crypto/ssh to prompt the user to verify host key
// fingerprint (same as OpenSSH does for an unknown host).
func HostKeyCallback(caCerts [][]byte) (ssh.HostKeyCallback, error) {
	var trustedKeys []ssh.PublicKey
	for _, caCert := range caCerts {
		for {
			_, _, publicKey, _, bytes, err := ssh.ParseKnownHosts(caCert)
			if err == io.EOF {
				break
			} else if err != nil {
				return nil, trace.Wrap(err, "failed parsing CA cert: %v; raw CA cert line: %q", err, caCert)
			}
			trustedKeys = append(trustedKeys, publicKey)
			caCert = bytes
		}
	}

	// No CAs are provided, return a nil callback which will prompt the user
	// for trust.
	if len(trustedKeys) == 0 {
		return nil, nil
	}

	return func(host string, a net.Addr, hostKey ssh.PublicKey) error {
		clusterCert, ok := hostKey.(*ssh.Certificate)
		if ok {
			hostKey = clusterCert.SignatureKey
		}
		for _, trustedKey := range trustedKeys {
			if KeysEqual(trustedKey, hostKey) {
				return nil
			}
		}
		return trace.AccessDenied("host %v is untrusted or Teleport CA has been rotated; try getting new credentials by logging in again ('tsh login') or re-exporting the identity file ('tctl auth sign' or 'tsh login -o'), depending on how you got them initially", host)
	}, nil
}

// KeysEqual is constant time compare of the keys to avoid timing attacks
func KeysEqual(ak, bk ssh.PublicKey) bool {
	a := ssh.Marshal(ak)
	b := ssh.Marshal(bk)
	return (len(a) == len(b) && subtle.ConstantTimeCompare(a, b) == 1)
}
<|endoftext|>


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

// +build appengine

package build

import (
	"bytes"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"strconv"

	"appengine"
	"appengine/datastore"
)

func init() {
	http.HandleFunc("/perf", perfChangesHandler)
}

// perfSummaryHandler draws the main benchmarking page.
func perfChangesHandler(w http.ResponseWriter, r *http.Request) {
	d := dashboardForRequest(r)
	c := d.Context(appengine.NewContext(r))

	page, _ := strconv.Atoi(r.FormValue("page"))
	if page < 0 {
		page = 0
	}

	pc, err := GetPerfConfig(c, r)
	if err != nil {
		logErr(w, r, err)
		return
	}

	commits, err := dashPerfCommits(c, page)
	if err != nil {
		logErr(w, r, err)
		return
	}

	// Fetch PerfResult's for the commits.
	var uiCommits []*perfChangesCommit
	rc := MakePerfResultCache(c, commits[0], false)

	// But first compare tip with the last release.
	if page == 0 {
		res0 := &PerfResult{CommitHash: knownTags[lastRelease]}
		if err := datastore.Get(c, res0.Key(c), res0); err != nil && err != datastore.ErrNoSuchEntity {
			logErr(w, r, fmt.Errorf("getting PerfResult: %v", err))
			return
		}
		if err != datastore.ErrNoSuchEntity {
			uiCom, err := handleOneCommit(pc, commits[0], rc, res0)
			if err != nil {
				logErr(w, r, err)
				return
			}
			uiCom.IsSummary = true
			uiCom.ParentHash = lastRelease
			uiCommits = append(uiCommits, uiCom)
		}
	}

	for _, com := range commits {
		uiCom, err := handleOneCommit(pc, com, rc, nil)
		if err != nil {
			logErr(w, r, err)
			return
		}
		uiCommits = append(uiCommits, uiCom)
	}

	p := &Pagination{}
	if len(commits) == commitsPerPage {
		p.Next = page + 1
	}
	if page > 0 {
		p.Prev = page - 1
		p.HasPrev = true
	}

	data := &perfChangesData{d, p, uiCommits}

	var buf bytes.Buffer
	if err := perfChangesTemplate.Execute(&buf, data); err != nil {
		logErr(w, r, err)
		return
	}

	buf.WriteTo(w)
}

func handleOneCommit(pc *PerfConfig, com *Commit, rc *PerfResultCache, baseRes *PerfResult) (*perfChangesCommit, error) {
	uiCom := new(perfChangesCommit)
	uiCom.Commit = com
	res1 := rc.Get(com.Num)
	for builder, benchmarks1 := range res1.ParseData() {
		for benchmark, data1 := range benchmarks1 {
			if benchmark != "meta-done" || !data1.OK {
				uiCom.NumResults++
			}
			if !data1.OK {
				v := new(perfChangesChange)
				v.diff = 10000
				v.Style = "fail"
				v.Builder = builder
				v.Link = fmt.Sprintf("log/%v", data1.Artifacts["log"])
				v.Val = builder
				v.Hint = builder
				if benchmark != "meta-done" {
					v.Hint += "/" + benchmark
				}
				m := findMetric(uiCom, "failure")
				m.BadChanges = append(m.BadChanges, v)
			}
		}
		res0 := baseRes
		if res0 == nil {
			var err error
			res0, err = rc.NextForComparison(com.Num, builder)
			if err != nil {
				return nil, err
			}
			if res0 == nil {
				continue
			}
		}
		changes := significantPerfChanges(pc, builder, res0, res1)
		for _, ch := range changes {
			v := new(perfChangesChange)
			v.Builder = builder
			v.Benchmark, v.Procs = splitBench(ch.Bench)
			v.diff = ch.Diff
			v.Val = fmt.Sprintf("%+.2f%%", ch.Diff)
			v.Hint = fmt.Sprintf("%v/%v", builder, ch.Bench)
			v.Link = fmt.Sprintf("perfdetail?commit=%v&commit0=%v&builder=%v&benchmark=%v", com.Hash, res0.CommitHash, builder, v.Benchmark)
			m := findMetric(uiCom, ch.Metric)
			if v.diff > 0 {
				v.Style = "bad"
				m.BadChanges = append(m.BadChanges, v)
			} else {
				v.Style = "good"
				m.GoodChanges = append(m.GoodChanges, v)
			}
		}
	}

	// Sort metrics and changes.
	for _, m := range uiCom.Metrics {
		sort.Sort(m.GoodChanges)
		sort.Sort(m.BadChanges)
	}
	sort.Sort(uiCom.Metrics)
	// Need at least one metric for UI.
	if len(uiCom.Metrics) == 0 {
		uiCom.Metrics = append(uiCom.Metrics, &perfChangesMetric{})
	}
	uiCom.Metrics[0].First = true
	return uiCom, nil
}

func findMetric(c *perfChangesCommit, metric string) *perfChangesMetric {
	for _, m := range c.Metrics {
		if m.Name == metric {
			return m
		}
	}
	m := new(perfChangesMetric)
	m.Name = metric
	c.Metrics = append(c.Metrics, m)
	return m
}

type uiPerfConfig struct {
	Builders   []uiPerfConfigElem
	Benchmarks []uiPerfConfigElem
	Metrics    []uiPerfConfigElem
	Procs      []uiPerfConfigElem
}

type uiPerfConfigElem struct {
	Name     string
	Selected bool
}

var perfChangesTemplate = template.Must(
	template.New("perf_changes.html").Funcs(tmplFuncs).ParseFiles("build/perf_changes.html"),
)

type perfChangesData struct {
	Dashboard  *Dashboard
	Pagination *Pagination
	Commits    []*perfChangesCommit
}

type perfChangesCommit struct {
	*Commit
	IsSummary  bool
	NumResults int
	Metrics    perfChangesMetricSlice
}

type perfChangesMetric struct {
	Name        string
	First       bool
	BadChanges  perfChangesChangeSlice
	GoodChanges perfChangesChangeSlice
}

type perfChangesChange struct {
	Builder   string
	Benchmark string
	Link      string
	Hint      string
	Style     string
	Val       string
	Procs     int
	diff      float64
}

type perfChangesMetricSlice []*perfChangesMetric

func (l perfChangesMetricSlice) Len() int      { return len(l) }
func (l perfChangesMetricSlice) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l perfChangesMetricSlice) Less(i, j int) bool {
	if l[i].Name == "failure" || l[j].Name == "failure" {
		return l[i].Name == "failure"
	}
	return l[i].Name < l[j].Name
}

type perfChangesChangeSlice []*perfChangesChange

func (l perfChangesChangeSlice) Len() int      { return len(l) }
func (l perfChangesChangeSlice) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l perfChangesChangeSlice) Less(i, j int) bool {
	vi, vj := l[i].diff, l[j].diff
	if vi > 0 && vj > 0 {
		return vi > vj
	} else if vi < 0 && vj < 0 {
		return vi < vj
	} else {
		panic("comparing positive and negative diff")
	}
}
<commit_msg>dashboard/app: dedup changes in UI Currently for every benchmark/metric we show all changes for all builders x procs. With 4 builders and 5 procs, that's 20 changes (20 red/green boxes in a single cell). Instead show only maximum change for every benchmark/metric. This significantly reduces clutter in UI. When you click on the red/green box, you can see the rest of the changes.<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build appengine

package build

import (
	"bytes"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"strconv"

	"appengine"
	"appengine/datastore"
)

func init() {
	http.HandleFunc("/perf", perfChangesHandler)
}

// perfSummaryHandler draws the main benchmarking page.
func perfChangesHandler(w http.ResponseWriter, r *http.Request) {
	d := dashboardForRequest(r)
	c := d.Context(appengine.NewContext(r))

	page, _ := strconv.Atoi(r.FormValue("page"))
	if page < 0 {
		page = 0
	}

	pc, err := GetPerfConfig(c, r)
	if err != nil {
		logErr(w, r, err)
		return
	}

	commits, err := dashPerfCommits(c, page)
	if err != nil {
		logErr(w, r, err)
		return
	}

	// Fetch PerfResult's for the commits.
	var uiCommits []*perfChangesCommit
	rc := MakePerfResultCache(c, commits[0], false)

	// But first compare tip with the last release.
	if page == 0 {
		res0 := &PerfResult{CommitHash: knownTags[lastRelease]}
		if err := datastore.Get(c, res0.Key(c), res0); err != nil && err != datastore.ErrNoSuchEntity {
			logErr(w, r, fmt.Errorf("getting PerfResult: %v", err))
			return
		}
		if err != datastore.ErrNoSuchEntity {
			uiCom, err := handleOneCommit(pc, commits[0], rc, res0)
			if err != nil {
				logErr(w, r, err)
				return
			}
			uiCom.IsSummary = true
			uiCom.ParentHash = lastRelease
			uiCommits = append(uiCommits, uiCom)
		}
	}

	for _, com := range commits {
		uiCom, err := handleOneCommit(pc, com, rc, nil)
		if err != nil {
			logErr(w, r, err)
			return
		}
		uiCommits = append(uiCommits, uiCom)
	}

	p := &Pagination{}
	if len(commits) == commitsPerPage {
		p.Next = page + 1
	}
	if page > 0 {
		p.Prev = page - 1
		p.HasPrev = true
	}

	data := &perfChangesData{d, p, uiCommits}

	var buf bytes.Buffer
	if err := perfChangesTemplate.Execute(&buf, data); err != nil {
		logErr(w, r, err)
		return
	}

	buf.WriteTo(w)
}

func handleOneCommit(pc *PerfConfig, com *Commit, rc *PerfResultCache, baseRes *PerfResult) (*perfChangesCommit, error) {
	uiCom := new(perfChangesCommit)
	uiCom.Commit = com
	res1 := rc.Get(com.Num)
	for builder, benchmarks1 := range res1.ParseData() {
		for benchmark, data1 := range benchmarks1 {
			if benchmark != "meta-done" || !data1.OK {
				uiCom.NumResults++
			}
			if !data1.OK {
				v := new(perfChangesChange)
				v.diff = 10000
				v.Style = "fail"
				v.Builder = builder
				v.Link = fmt.Sprintf("log/%v", data1.Artifacts["log"])
				v.Val = builder
				v.Hint = builder
				if benchmark != "meta-done" {
					v.Hint += "/" + benchmark
				}
				m := findMetric(uiCom, "failure")
				m.BadChanges = append(m.BadChanges, v)
			}
		}
		res0 := baseRes
		if res0 == nil {
			var err error
			res0, err = rc.NextForComparison(com.Num, builder)
			if err != nil {
				return nil, err
			}
			if res0 == nil {
				continue
			}
		}
		changes := significantPerfChanges(pc, builder, res0, res1)
		changes = dedupPerfChanges(changes)
		for _, ch := range changes {
			v := new(perfChangesChange)
			v.Builder = builder
			v.Benchmark, v.Procs = splitBench(ch.Bench)
			v.diff = ch.Diff
			v.Val = fmt.Sprintf("%+.2f%%", ch.Diff)
			v.Hint = fmt.Sprintf("%v/%v", builder, ch.Bench)
			v.Link = fmt.Sprintf("perfdetail?commit=%v&commit0=%v&builder=%v&benchmark=%v", com.Hash, res0.CommitHash, builder, v.Benchmark)
			m := findMetric(uiCom, ch.Metric)
			if v.diff > 0 {
				v.Style = "bad"
				m.BadChanges = append(m.BadChanges, v)
			} else {
				v.Style = "good"
				m.GoodChanges = append(m.GoodChanges, v)
			}
		}
	}

	// Sort metrics and changes.
	for _, m := range uiCom.Metrics {
		sort.Sort(m.GoodChanges)
		sort.Sort(m.BadChanges)
	}
	sort.Sort(uiCom.Metrics)
	// Need at least one metric for UI.
	if len(uiCom.Metrics) == 0 {
		uiCom.Metrics = append(uiCom.Metrics, &perfChangesMetric{})
	}
	uiCom.Metrics[0].First = true
	return uiCom, nil
}

// Find builder-procs with the maximum absolute diff for every benchmark-metric, drop the rest.
func dedupPerfChanges(changes []*PerfChange) (deduped []*PerfChange) {
	maxDiff := make(map[string]float64)
	maxBench := make(map[string]string)
	// First, find the maximum.
	for _, ch := range changes {
		bench, _ := splitBench(ch.Bench)
		k := bench + "|" + ch.Metric
		v := ch.Diff
		if v < 0 {
			v = -v
		}
		if maxDiff[k] < v {
			maxDiff[k] = v
			maxBench[k] = ch.Builder + "|" + ch.Bench
		}
	}
	// Then, remove the rest.
	for _, ch := range changes {
		bench, _ := splitBench(ch.Bench)
		k := bench + "|" + ch.Metric
		if maxBench[k] == ch.Builder+"|"+ch.Bench {
			deduped = append(deduped, ch)
		}
	}
	return
}

func findMetric(c *perfChangesCommit, metric string) *perfChangesMetric {
	for _, m := range c.Metrics {
		if m.Name == metric {
			return m
		}
	}
	m := new(perfChangesMetric)
	m.Name = metric
	c.Metrics = append(c.Metrics, m)
	return m
}

type uiPerfConfig struct {
	Builders   []uiPerfConfigElem
	Benchmarks []uiPerfConfigElem
	Metrics    []uiPerfConfigElem
	Procs      []uiPerfConfigElem
}

type uiPerfConfigElem struct {
	Name     string
	Selected bool
}

var perfChangesTemplate = template.Must(
	template.New("perf_changes.html").Funcs(tmplFuncs).ParseFiles("build/perf_changes.html"),
)

type perfChangesData struct {
	Dashboard  *Dashboard
	Pagination *Pagination
	Commits    []*perfChangesCommit
}

type perfChangesCommit struct {
	*Commit
	IsSummary  bool
	NumResults int
	Metrics    perfChangesMetricSlice
}

type perfChangesMetric struct {
	Name        string
	First       bool
	BadChanges  perfChangesChangeSlice
	GoodChanges perfChangesChangeSlice
}

type perfChangesChange struct {
	Builder   string
	Benchmark string
	Link      string
	Hint      string
	Style     string
	Val       string
	Procs     int
	diff      float64
}

type perfChangesMetricSlice []*perfChangesMetric

func (l perfChangesMetricSlice) Len() int      { return len(l) }
func (l perfChangesMetricSlice) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l perfChangesMetricSlice) Less(i, j int) bool {
	if l[i].Name == "failure" || l[j].Name == "failure" {
		return l[i].Name == "failure"
	}
	return l[i].Name < l[j].Name
}

type perfChangesChangeSlice []*perfChangesChange

func (l perfChangesChangeSlice) Len() int      { return len(l) }
func (l perfChangesChangeSlice) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l perfChangesChangeSlice) Less(i, j int) bool {
	vi, vj := l[i].diff, l[j].diff
	if vi > 0 && vj > 0 {
		return vi > vj
	} else if vi < 0 && vj < 0 {
		return vi < vj
	} else {
		panic("comparing positive and negative diff")
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Ulrich Kunitz. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The directory basics provides packages for standard types providing
// basic functions as addition and subtraction with overflow detection.
package basics
<commit_msg>basics: updated documentation<commit_after>// Copyright 2015 Ulrich Kunitz. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The directory basics contains packages providing basic functions for
// standard types. Example for those functions are addition and
// subtraction with overflow detection.
package basics
<|endoftext|>


<commit_before>package metric

import (
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"

	flags "github.com/jessevdk/go-flags"
)

type Event struct {
	Name       string
	Value      interface{}
	State      EventState
	Attributes map[string]string
	Host       string
	Time       time.Time
}

type EventState string

const (
	EventStateOK       EventState = "ok"
	EventStateWarning  EventState = "warning"
	EventStateCritical EventState = "critical"
)

//go:generate counterfeiter . Emitter
type Emitter interface {
	Emit(lager.Logger, Event)
}

//go:generate counterfeiter . EmitterFactory
type EmitterFactory interface {
	Description() string
	IsConfigured() bool
	NewEmitter() (Emitter, error)
}

var emitterFactories []EmitterFactory

func RegisterEmitter(factory EmitterFactory) {
	emitterFactories = append(emitterFactories, factory)
}

func WireEmitters(group *flags.Group) {
	for _, factory := range emitterFactories {
		_, err := group.AddGroup(fmt.Sprintf("Metric Emitter (%s)", factory.Description()), "", factory)
		if err != nil {
			panic(err)
		}
	}
}

type eventEmission struct {
	event  Event
	logger lager.Logger
}

var (
	emitter         Emitter
	eventHost       string
	eventAttributes map[string]string
	emissions       chan eventEmission
)

func Initialize(logger lager.Logger, host string, attributes map[string]string, bufferSize uint32) error {
	logger.Debug("metric-initialize", lager.Data{
		"host": host,
		"attributes": attributes,
		"buffer-size": bufferSize,
	})

	var (
		emitterDescriptions []string
		err                 error
	)

	for _, factory := range emitterFactories {
		if factory.IsConfigured() {
			emitterDescriptions = append(emitterDescriptions, factory.Description())
		}
	}
	if len(emitterDescriptions) > 1 {
		return fmt.Errorf("Multiple emitters configured: %s", strings.Join(emitterDescriptions, ", "))
	}

	for _, factory := range emitterFactories {
		if factory.IsConfigured() {
			emitter, err = factory.NewEmitter()
			if err != nil {
				return err
			}
		}
	}

	if emitter == nil {
		return nil
	}

	emitter = emitter
	eventHost = host
	eventAttributes = attributes
	emissions = make(chan eventEmission, int(bufferSize))

	go emitLoop()

	return nil
}

func Deinitialize(logger lager.Logger) {
	close(emissions)
	emitterFactories = nil
}

func emit(logger lager.Logger, event Event) {
	logger.Debug("emit-event", lager.Data{
		"event": event, "emitter": emitter,
	})
	if emitter == nil {
		return
	}

	event.Host = eventHost
	event.Time = time.Now()

	mergedAttributes := map[string]string{}
	for k, v := range eventAttributes {
		mergedAttributes[k] = v
	}

	if event.Attributes != nil {
		for k, v := range event.Attributes {
			mergedAttributes[k] = v
		}
	}

	event.Attributes = mergedAttributes

	select {
	case emissions <- eventEmission{logger: logger, event: event}:
		logger.Debug("emit-event-write-to-channel", lager.Data{
			"event": event,
		})
	default:
		logger.Error("queue-full", nil)
	}
}

func emitLoop() {
	for emission := range emissions {
		emission.logger.Debug("emit-event-loop", lager.Data{
			"event": emission.event,
		})
		emitter.Emit(emission.logger.Session("emit"), emission.event)
	}
}
<commit_msg>Remove extraneous logging<commit_after>package metric

import (
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"

	flags "github.com/jessevdk/go-flags"
)

type Event struct {
	Name       string
	Value      interface{}
	State      EventState
	Attributes map[string]string
	Host       string
	Time       time.Time
}

type EventState string

const (
	EventStateOK       EventState = "ok"
	EventStateWarning  EventState = "warning"
	EventStateCritical EventState = "critical"
)

//go:generate counterfeiter . Emitter
type Emitter interface {
	Emit(lager.Logger, Event)
}

//go:generate counterfeiter . EmitterFactory
type EmitterFactory interface {
	Description() string
	IsConfigured() bool
	NewEmitter() (Emitter, error)
}

var emitterFactories []EmitterFactory

func RegisterEmitter(factory EmitterFactory) {
	emitterFactories = append(emitterFactories, factory)
}

func WireEmitters(group *flags.Group) {
	for _, factory := range emitterFactories {
		_, err := group.AddGroup(fmt.Sprintf("Metric Emitter (%s)", factory.Description()), "", factory)
		if err != nil {
			panic(err)
		}
	}
}

type eventEmission struct {
	event  Event
	logger lager.Logger
}

var (
	emitter         Emitter
	eventHost       string
	eventAttributes map[string]string
	emissions       chan eventEmission
)

func Initialize(logger lager.Logger, host string, attributes map[string]string, bufferSize uint32) error {
	logger.Debug("metric-initialize", lager.Data{
		"host": host,
		"attributes": attributes,
		"buffer-size": bufferSize,
	})

	var (
		emitterDescriptions []string
		err                 error
	)

	for _, factory := range emitterFactories {
		if factory.IsConfigured() {
			emitterDescriptions = append(emitterDescriptions, factory.Description())
		}
	}
	if len(emitterDescriptions) > 1 {
		return fmt.Errorf("Multiple emitters configured: %s", strings.Join(emitterDescriptions, ", "))
	}

	for _, factory := range emitterFactories {
		if factory.IsConfigured() {
			emitter, err = factory.NewEmitter()
			if err != nil {
				return err
			}
		}
	}

	if emitter == nil {
		return nil
	}

	emitter = emitter
	eventHost = host
	eventAttributes = attributes
	emissions = make(chan eventEmission, int(bufferSize))

	go emitLoop()

	return nil
}

func Deinitialize(logger lager.Logger) {
	close(emissions)
	emitterFactories = nil
}

func emit(logger lager.Logger, event Event) {
	if emitter == nil {
		return
	}

	event.Host = eventHost
	event.Time = time.Now()

	mergedAttributes := map[string]string{}
	for k, v := range eventAttributes {
		mergedAttributes[k] = v
	}

	if event.Attributes != nil {
		for k, v := range event.Attributes {
			mergedAttributes[k] = v
		}
	}

	event.Attributes = mergedAttributes

	select {
	case emissions <- eventEmission{logger: logger, event: event}:
	default:
		logger.Error("queue-full", nil)
	}
}

func emitLoop() {
	for emission := range emissions {
		emitter.Emit(emission.logger.Session("emit"), emission.event)
	}
}
<|endoftext|>


<commit_before>package attacks

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/sourcekris/goRsaTool/attacks/brokenrsa"
	"github.com/sourcekris/goRsaTool/attacks/commonfactor"
	"github.com/sourcekris/goRsaTool/attacks/commonmodulus"
	"github.com/sourcekris/goRsaTool/attacks/crt"
	"github.com/sourcekris/goRsaTool/attacks/defectivee"
	"github.com/sourcekris/goRsaTool/attacks/factordb"
	"github.com/sourcekris/goRsaTool/attacks/fermat"
	"github.com/sourcekris/goRsaTool/attacks/franklinreiter"
	"github.com/sourcekris/goRsaTool/attacks/gmpecm"
	"github.com/sourcekris/goRsaTool/attacks/hastads"
	"github.com/sourcekris/goRsaTool/attacks/hastadsbroadcast"
	"github.com/sourcekris/goRsaTool/attacks/knownprime"
	"github.com/sourcekris/goRsaTool/attacks/londahl"
	"github.com/sourcekris/goRsaTool/attacks/manysmallprimes"
	"github.com/sourcekris/goRsaTool/attacks/notableprimes"
	"github.com/sourcekris/goRsaTool/attacks/oraclemodulus"
	"github.com/sourcekris/goRsaTool/attacks/partiald"
	"github.com/sourcekris/goRsaTool/attacks/pastctfprimes"
	"github.com/sourcekris/goRsaTool/attacks/pollardrhobrent"
	"github.com/sourcekris/goRsaTool/attacks/pollardsp1"
	"github.com/sourcekris/goRsaTool/attacks/pollardsrho"
	"github.com/sourcekris/goRsaTool/attacks/qicheng"
	"github.com/sourcekris/goRsaTool/attacks/smallfractions"
	"github.com/sourcekris/goRsaTool/attacks/smallq"
	"github.com/sourcekris/goRsaTool/attacks/squaren"
	"github.com/sourcekris/goRsaTool/attacks/wiener"
	"github.com/sourcekris/goRsaTool/attacks/wienermultiprime"
	"github.com/sourcekris/goRsaTool/attacks/williamsp1"
	"github.com/sourcekris/goRsaTool/keys"
)

// default timeout 3m0s
const DefaultTimeout int = 180

// SupportedAttacks stores the list of registered attacks we support.
var SupportedAttacks *Attacks

func init() {
	SupportedAttacks = NewAttacks()
	// TODO(sourcekris): Register them in each package init function.
	SupportedAttacks.RegisterAttack("crtsolver", false, true, DefaultTimeout, crt.Attack)
	SupportedAttacks.RegisterAttack("factordb", false, true, DefaultTimeout, factordb.Attack)
	SupportedAttacks.RegisterAttack("hastads", false, true, DefaultTimeout, hastads.Attack)
	SupportedAttacks.RegisterAttack("hastadsbroadcast", true, true, DefaultTimeout, hastadsbroadcast.Attack)
	SupportedAttacks.RegisterAttack("commonfactors", true, true, DefaultTimeout, commonfactor.Attack)
	SupportedAttacks.RegisterAttack("commonmodulus", true, true, DefaultTimeout, commonmodulus.Attack)
	SupportedAttacks.RegisterAttack("partiald", false, false, DefaultTimeout, partiald.Attack)
	SupportedAttacks.RegisterAttack("knownprime", false, false, DefaultTimeout, knownprime.Attack)
	SupportedAttacks.RegisterAttack("brokenrsa", false, true, DefaultTimeout, brokenrsa.Attack)
	SupportedAttacks.RegisterAttack("notableprimes", false, true, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("pastctf", false, true, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("smallq", false, true, DefaultTimeout, smallq.Attack)
	SupportedAttacks.RegisterAttack("wiener", false, true, DefaultTimeout, wiener.Attack)
	SupportedAttacks.RegisterAttack("wienermultiprime", false, true, DefaultTimeout, wienermultiprime.Attack)
	SupportedAttacks.RegisterAttack("qicheng", false, true, DefaultTimeout, qicheng.Attack)
	SupportedAttacks.RegisterAttack("fermat", false, true, DefaultTimeout, fermat.Attack)
	SupportedAttacks.RegisterAttack("londahl", false, true, DefaultTimeout, londahl.Attack)
	SupportedAttacks.RegisterAttack("smallfractions", false, true, DefaultTimeout, smallfractions.Attack)
	SupportedAttacks.RegisterAttack("manysmallprimes", false, true, DefaultTimeout, manysmallprimes.Attack)
	SupportedAttacks.RegisterAttack("ecm", false, true, DefaultTimeout, gmpecm.Attack)
	SupportedAttacks.RegisterAttack("franklinreiter", true, true, DefaultTimeout, franklinreiter.Attack)
	SupportedAttacks.RegisterAttack("pollardsp1", false, true, DefaultTimeout, pollardsp1.Attack)
	SupportedAttacks.RegisterAttack("pollardsrho", false, true, DefaultTimeout, pollardsrho.Attack)
	SupportedAttacks.RegisterAttack("pollardrhobrent", false, true, DefaultTimeout, pollardrhobrent.Attack)
	SupportedAttacks.RegisterAttack("williamsp1", false, true, DefaultTimeout, williamsp1.Attack)
	SupportedAttacks.RegisterAttack("defectivee", false, true, DefaultTimeout, defectivee.Attack)
	SupportedAttacks.RegisterAttack("oraclemodulus", false, true, DefaultTimeout, oraclemodulus.Attack)
	SupportedAttacks.RegisterAttack("squaren", false, true, DefaultTimeout, squaren.Attack)

	// Aliased attacks (names that point to attacks already in the above list).
	SupportedAttacks.RegisterAttack("mersenne", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("lucas", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("novelty", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("pastprimes", false, false, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("pastctfprimes", false, false, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("sexyprimes", false, false, DefaultTimeout, fermat.Attack)
}

type attackFunc func([]*keys.RSA, chan error)

// Attack encodes a single attack and what features it supports.
type Attack struct {
	Name          string
	SupportsMulti bool
	Unnatended    bool
	Timeout       int
	F             attackFunc
}

// Attacks wraps a slice of Attack objects that are supported.
type Attacks struct {
	Supported []*Attack
}

// NewAttacks constructs a new Attacks object.
func NewAttacks() *Attacks {
	return &Attacks{}
}

// RegisterAttack adds a new attack to the receiving Attacks.
func (a *Attacks) RegisterAttack(name string, multi bool, unnatended bool, timeout int, f attackFunc) {
	if a == nil {
		a = NewAttacks()
	}

	a.Supported = append(a.Supported, &Attack{name, multi, unnatended, timeout, f})
}

// IsSupported returns true if name attack is supported.
func (a *Attacks) IsSupported(name string) bool {
	for _, a := range a.Supported {
		if a.Name == name {
			return true
		}
	}

	return false
}

// SupportsMulti returns true if the attack supports multi-key attacks.
func (a *Attacks) SupportsMulti(name string) bool {
	for _, a := range a.Supported {
		if a.Name == name {
			return a.SupportsMulti
		}
	}

	return false
}

// Execute executes the named attack against t.
func (a *Attacks) Execute(name string, t []*keys.RSA) error {
	if SupportedAttacks == nil {
		return errors.New("no attacks registered")
	}

	if !a.IsSupported(name) {
		return fmt.Errorf("unsupported attack: %v", name)
	}

	for _, a := range SupportedAttacks.Supported {
		if a.Name == name {
			ctx := context.Background()
			ctx, cancel := context.WithTimeout(ctx, time.Duration(a.Timeout*int(time.Second)))
			defer cancel()

			ch := make(chan error)

			go a.F(t, ch)

			select {
			case result := <-ch:
				return result
			case <-ctx.Done():
				return fmt.Errorf("%s failed to factorize the key in the given time", name)
			}
		}
	}

	return errors.New("attack not found")
}
<commit_msg>change timeout for pollardsrho and pollardsrhobrent<commit_after>package attacks

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/sourcekris/goRsaTool/attacks/brokenrsa"
	"github.com/sourcekris/goRsaTool/attacks/commonfactor"
	"github.com/sourcekris/goRsaTool/attacks/commonmodulus"
	"github.com/sourcekris/goRsaTool/attacks/crt"
	"github.com/sourcekris/goRsaTool/attacks/defectivee"
	"github.com/sourcekris/goRsaTool/attacks/factordb"
	"github.com/sourcekris/goRsaTool/attacks/fermat"
	"github.com/sourcekris/goRsaTool/attacks/franklinreiter"
	"github.com/sourcekris/goRsaTool/attacks/gmpecm"
	"github.com/sourcekris/goRsaTool/attacks/hastads"
	"github.com/sourcekris/goRsaTool/attacks/hastadsbroadcast"
	"github.com/sourcekris/goRsaTool/attacks/knownprime"
	"github.com/sourcekris/goRsaTool/attacks/londahl"
	"github.com/sourcekris/goRsaTool/attacks/manysmallprimes"
	"github.com/sourcekris/goRsaTool/attacks/notableprimes"
	"github.com/sourcekris/goRsaTool/attacks/oraclemodulus"
	"github.com/sourcekris/goRsaTool/attacks/partiald"
	"github.com/sourcekris/goRsaTool/attacks/pastctfprimes"
	"github.com/sourcekris/goRsaTool/attacks/pollardrhobrent"
	"github.com/sourcekris/goRsaTool/attacks/pollardsp1"
	"github.com/sourcekris/goRsaTool/attacks/pollardsrho"
	"github.com/sourcekris/goRsaTool/attacks/qicheng"
	"github.com/sourcekris/goRsaTool/attacks/smallfractions"
	"github.com/sourcekris/goRsaTool/attacks/smallq"
	"github.com/sourcekris/goRsaTool/attacks/squaren"
	"github.com/sourcekris/goRsaTool/attacks/wiener"
	"github.com/sourcekris/goRsaTool/attacks/wienermultiprime"
	"github.com/sourcekris/goRsaTool/attacks/williamsp1"
	"github.com/sourcekris/goRsaTool/keys"
)

// default timeout 3m0s
const DefaultTimeout int = 180

// SupportedAttacks stores the list of registered attacks we support.
var SupportedAttacks *Attacks

func init() {
	SupportedAttacks = NewAttacks()
	// TODO(sourcekris): Register them in each package init function.
	SupportedAttacks.RegisterAttack("crtsolver", false, true, DefaultTimeout, crt.Attack)
	SupportedAttacks.RegisterAttack("factordb", false, true, DefaultTimeout, factordb.Attack)
	SupportedAttacks.RegisterAttack("hastads", false, true, DefaultTimeout, hastads.Attack)
	SupportedAttacks.RegisterAttack("hastadsbroadcast", true, true, DefaultTimeout, hastadsbroadcast.Attack)
	SupportedAttacks.RegisterAttack("commonfactors", true, true, DefaultTimeout, commonfactor.Attack)
	SupportedAttacks.RegisterAttack("commonmodulus", true, true, DefaultTimeout, commonmodulus.Attack)
	SupportedAttacks.RegisterAttack("partiald", false, false, DefaultTimeout, partiald.Attack)
	SupportedAttacks.RegisterAttack("knownprime", false, false, DefaultTimeout, knownprime.Attack)
	SupportedAttacks.RegisterAttack("brokenrsa", false, true, DefaultTimeout, brokenrsa.Attack)
	SupportedAttacks.RegisterAttack("notableprimes", false, true, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("pastctf", false, true, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("smallq", false, true, DefaultTimeout, smallq.Attack)
	SupportedAttacks.RegisterAttack("wiener", false, true, DefaultTimeout, wiener.Attack)
	SupportedAttacks.RegisterAttack("wienermultiprime", false, true, DefaultTimeout, wienermultiprime.Attack)
	SupportedAttacks.RegisterAttack("qicheng", false, true, DefaultTimeout, qicheng.Attack)
	SupportedAttacks.RegisterAttack("fermat", false, true, DefaultTimeout, fermat.Attack)
	SupportedAttacks.RegisterAttack("londahl", false, true, DefaultTimeout, londahl.Attack)
	SupportedAttacks.RegisterAttack("smallfractions", false, true, DefaultTimeout, smallfractions.Attack)
	SupportedAttacks.RegisterAttack("manysmallprimes", false, true, DefaultTimeout, manysmallprimes.Attack)
	SupportedAttacks.RegisterAttack("ecm", false, true, DefaultTimeout, gmpecm.Attack)
	SupportedAttacks.RegisterAttack("franklinreiter", true, true, DefaultTimeout, franklinreiter.Attack)
	SupportedAttacks.RegisterAttack("pollardsp1", false, true, DefaultTimeout, pollardsp1.Attack)
	SupportedAttacks.RegisterAttack("pollardsrho", false, true, 300, pollardsrho.Attack)
	SupportedAttacks.RegisterAttack("pollardrhobrent", false, true, 300, pollardrhobrent.Attack)
	SupportedAttacks.RegisterAttack("williamsp1", false, true, DefaultTimeout, williamsp1.Attack)
	SupportedAttacks.RegisterAttack("defectivee", false, true, DefaultTimeout, defectivee.Attack)
	SupportedAttacks.RegisterAttack("oraclemodulus", false, true, DefaultTimeout, oraclemodulus.Attack)
	SupportedAttacks.RegisterAttack("squaren", false, true, DefaultTimeout, squaren.Attack)

	// Aliased attacks (names that point to attacks already in the above list).
	SupportedAttacks.RegisterAttack("mersenne", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("lucas", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("novelty", false, false, DefaultTimeout, notableprimes.Attack)
	SupportedAttacks.RegisterAttack("pastprimes", false, false, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("pastctfprimes", false, false, DefaultTimeout, pastctfprimes.Attack)
	SupportedAttacks.RegisterAttack("sexyprimes", false, false, DefaultTimeout, fermat.Attack)
}

type attackFunc func([]*keys.RSA, chan error)

// Attack encodes a single attack and what features it supports.
type Attack struct {
	Name          string
	SupportsMulti bool
	Unnatended    bool
	Timeout       int
	F             attackFunc
}

// Attacks wraps a slice of Attack objects that are supported.
type Attacks struct {
	Supported []*Attack
}

// NewAttacks constructs a new Attacks object.
func NewAttacks() *Attacks {
	return &Attacks{}
}

// RegisterAttack adds a new attack to the receiving Attacks.
func (a *Attacks) RegisterAttack(name string, multi bool, unnatended bool, timeout int, f attackFunc) {
	if a == nil {
		a = NewAttacks()
	}

	a.Supported = append(a.Supported, &Attack{name, multi, unnatended, timeout, f})
}

// IsSupported returns true if name attack is supported.
func (a *Attacks) IsSupported(name string) bool {
	for _, a := range a.Supported {
		if a.Name == name {
			return true
		}
	}

	return false
}

// SupportsMulti returns true if the attack supports multi-key attacks.
func (a *Attacks) SupportsMulti(name string) bool {
	for _, a := range a.Supported {
		if a.Name == name {
			return a.SupportsMulti
		}
	}

	return false
}

// Execute executes the named attack against t.
func (a *Attacks) Execute(name string, t []*keys.RSA) error {
	if SupportedAttacks == nil {
		return errors.New("no attacks registered")
	}

	if !a.IsSupported(name) {
		return fmt.Errorf("unsupported attack: %v", name)
	}

	for _, a := range SupportedAttacks.Supported {
		if a.Name == name {
			ctx := context.Background()
			ctx, cancel := context.WithTimeout(ctx, time.Duration(a.Timeout*int(time.Second)))
			defer cancel()

			ch := make(chan error)

			go a.F(t, ch)

			select {
			case result := <-ch:
				return result
			case <-ctx.Done():
				return fmt.Errorf("%s failed to factorize the key in the given time", name)
			}
		}
	}

	return errors.New("attack not found")
}
<|endoftext|>