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

package oauth

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/auth"
	"github.com/tsuru/tsuru/auth/native"
	tsuruErrors "github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/log"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

var (
	ErrMissingCodeError       = &tsuruErrors.ValidationError{Message: "You must provide code to login"}
	ErrMissingCodeRedirectUrl = &tsuruErrors.ValidationError{Message: "You must provide the used redirect url to login"}
	ErrEmptyAccessToken       = &tsuruErrors.NotAuthorizedError{Message: "Couldn't convert code to access token."}
	ErrEmptyUserEmail         = &tsuruErrors.NotAuthorizedError{Message: "Couldn't parse user email."}

	requestLatencies = prometheus.NewHistogram(prometheus.HistogramOpts{
		Name: "tsuru_oath_request_duration_seconds",
		Help: "The oath requests latency distributions.",
	})
	requestErrors = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "tsuru_oath_request_errors_total",
		Help: "The total number of oath request errors.",
	})
)

type OAuthParser interface {
	Parse(infoResponse *http.Response) (string, error)
}

type OAuthScheme struct {
	BaseConfig   oauth2.Config
	InfoUrl      string
	CallbackPort int
	Parser       OAuthParser
}

func init() {
	auth.RegisterScheme("oauth", &OAuthScheme{})
	prometheus.MustRegister(requestLatencies)
	prometheus.MustRegister(requestErrors)
}

// This method loads basic config and returns a copy of the
// config object.
func (s *OAuthScheme) loadConfig() (oauth2.Config, error) {
	if s.BaseConfig.ClientID != "" {
		return s.BaseConfig, nil
	}
	if s.Parser == nil {
		s.Parser = s
	}
	var emptyConfig oauth2.Config
	clientId, err := config.GetString("auth:oauth:client-id")
	if err != nil {
		return emptyConfig, err
	}
	clientSecret, err := config.GetString("auth:oauth:client-secret")
	if err != nil {
		return emptyConfig, err
	}
	scope, err := config.GetString("auth:oauth:scope")
	if err != nil {
		return emptyConfig, err
	}
	authURL, err := config.GetString("auth:oauth:auth-url")
	if err != nil {
		return emptyConfig, err
	}
	tokenURL, err := config.GetString("auth:oauth:token-url")
	if err != nil {
		return emptyConfig, err
	}
	infoURL, err := config.GetString("auth:oauth:info-url")
	if err != nil {
		return emptyConfig, err
	}
	callbackPort, err := config.GetInt("auth:oauth:callback-port")
	if err != nil {
		log.Debugf("auth:oauth:callback-port not found using random port: %s", err)
	}
	s.InfoUrl = infoURL
	s.CallbackPort = callbackPort
	s.BaseConfig = oauth2.Config{
		ClientID:     clientId,
		ClientSecret: clientSecret,
		Scopes:       []string{scope},
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
	}
	return s.BaseConfig, nil
}

func (s *OAuthScheme) Login(params map[string]string) (auth.Token, error) {
	conf, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	code, ok := params["code"]
	if !ok {
		return nil, ErrMissingCodeError
	}
	redirectUrl, ok := params["redirectUrl"]
	if !ok {
		return nil, ErrMissingCodeRedirectUrl
	}
	conf.RedirectURL = redirectUrl
	oauthToken, err := conf.Exchange(context.Background(), code)
	if err != nil {
		return nil, err
	}
	return s.handleToken(oauthToken)
}

func (s *OAuthScheme) handleToken(t *oauth2.Token) (*Token, error) {
	if t.AccessToken == "" {
		return nil, ErrEmptyAccessToken
	}
	conf, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	client := conf.Client(context.Background(), t)
	t0 := time.Now()
	response, err := client.Get(s.InfoUrl)
	requestLatencies.Observe(float64(time.Since(t0).Seconds()))
	if err != nil {
		requestErrors.Inc()
		return nil, err
	}
	defer response.Body.Close()
	email, err := s.Parser.Parse(response)
	if err != nil {
		return nil, err
	}
	if email == "" {
		return nil, ErrEmptyUserEmail
	}
	user, err := auth.GetUserByEmail(email)
	if err != nil {
		if err != auth.ErrUserNotFound {
			return nil, err
		}
		registrationEnabled, _ := config.GetBool("auth:user-registration")
		if !registrationEnabled {
			return nil, err
		}
		user = &auth.User{Email: email}
		err = user.Create()
		if err != nil {
			return nil, err
		}
	}
	token := Token{*t, email}
	err = token.save()
	if err != nil {
		return nil, err
	}
	return &token, nil
}

func (s *OAuthScheme) AppLogin(appName string) (auth.Token, error) {
	nativeScheme := native.NativeScheme{}
	return nativeScheme.AppLogin(appName)
}

func (s *OAuthScheme) AppLogout(token string) error {
	nativeScheme := native.NativeScheme{}
	return nativeScheme.AppLogout(token)
}

func (s *OAuthScheme) Logout(token string) error {
	return deleteToken(token)
}

func (s *OAuthScheme) Auth(header string) (auth.Token, error) {
	token, err := getToken(header)
	if err != nil {
		nativeScheme := native.NativeScheme{}
		token, nativeErr := nativeScheme.Auth(header)
		if nativeErr == nil && token.IsAppToken() {
			return token, nil
		}
		return nil, err
	}
	config, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	client := config.Client(context.Background(), &token.Token)
	t0 := time.Now()
	rsp, err := client.Get(s.InfoUrl)
	requestLatencies.Observe(float64(time.Since(t0).Seconds()))
	if err != nil {
		requestErrors.Inc()
		return nil, err
	}
	defer rsp.Body.Close()
	return token, nil
}

func (s *OAuthScheme) Name() string {
	return "oauth"
}

func (s *OAuthScheme) Info() (auth.SchemeInfo, error) {
	config, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	config.RedirectURL = "__redirect_url__"
	return auth.SchemeInfo{"authorizeUrl": config.AuthCodeURL(""), "port": strconv.Itoa(s.CallbackPort)}, nil
}

func (s *OAuthScheme) Parse(infoResponse *http.Response) (string, error) {
	user := struct {
		Email string `json:"email"`
	}{}
	data, err := ioutil.ReadAll(infoResponse.Body)
	if err != nil {
		return "", errors.Wrap(err, "unable to read user data response")
	}
	if infoResponse.StatusCode != http.StatusOK {
		return "", errors.Errorf("unexpected user data response %d: %s", infoResponse.StatusCode, data)
	}
	err = json.Unmarshal(data, &user)
	if err != nil {
		return "", errors.Wrapf(err, "unable to parse user data: %s", data)
	}
	return user.Email, nil
}

func (s *OAuthScheme) Create(user *auth.User) (*auth.User, error) {
	user.Password = ""
	err := user.Create()
	if err != nil {
		return nil, err
	}
	return user, nil
}

func (s *OAuthScheme) Remove(u *auth.User) error {
	err := deleteAllTokens(u.Email)
	if err != nil {
		return err
	}
	return u.Delete()
}
<commit_msg>auth: remove unnecesary conversion<commit_after>// Copyright 2017 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package oauth

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/auth"
	"github.com/tsuru/tsuru/auth/native"
	tsuruErrors "github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/log"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

var (
	ErrMissingCodeError       = &tsuruErrors.ValidationError{Message: "You must provide code to login"}
	ErrMissingCodeRedirectUrl = &tsuruErrors.ValidationError{Message: "You must provide the used redirect url to login"}
	ErrEmptyAccessToken       = &tsuruErrors.NotAuthorizedError{Message: "Couldn't convert code to access token."}
	ErrEmptyUserEmail         = &tsuruErrors.NotAuthorizedError{Message: "Couldn't parse user email."}

	requestLatencies = prometheus.NewHistogram(prometheus.HistogramOpts{
		Name: "tsuru_oath_request_duration_seconds",
		Help: "The oath requests latency distributions.",
	})
	requestErrors = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "tsuru_oath_request_errors_total",
		Help: "The total number of oath request errors.",
	})
)

type OAuthParser interface {
	Parse(infoResponse *http.Response) (string, error)
}

type OAuthScheme struct {
	BaseConfig   oauth2.Config
	InfoUrl      string
	CallbackPort int
	Parser       OAuthParser
}

func init() {
	auth.RegisterScheme("oauth", &OAuthScheme{})
	prometheus.MustRegister(requestLatencies)
	prometheus.MustRegister(requestErrors)
}

// This method loads basic config and returns a copy of the
// config object.
func (s *OAuthScheme) loadConfig() (oauth2.Config, error) {
	if s.BaseConfig.ClientID != "" {
		return s.BaseConfig, nil
	}
	if s.Parser == nil {
		s.Parser = s
	}
	var emptyConfig oauth2.Config
	clientId, err := config.GetString("auth:oauth:client-id")
	if err != nil {
		return emptyConfig, err
	}
	clientSecret, err := config.GetString("auth:oauth:client-secret")
	if err != nil {
		return emptyConfig, err
	}
	scope, err := config.GetString("auth:oauth:scope")
	if err != nil {
		return emptyConfig, err
	}
	authURL, err := config.GetString("auth:oauth:auth-url")
	if err != nil {
		return emptyConfig, err
	}
	tokenURL, err := config.GetString("auth:oauth:token-url")
	if err != nil {
		return emptyConfig, err
	}
	infoURL, err := config.GetString("auth:oauth:info-url")
	if err != nil {
		return emptyConfig, err
	}
	callbackPort, err := config.GetInt("auth:oauth:callback-port")
	if err != nil {
		log.Debugf("auth:oauth:callback-port not found using random port: %s", err)
	}
	s.InfoUrl = infoURL
	s.CallbackPort = callbackPort
	s.BaseConfig = oauth2.Config{
		ClientID:     clientId,
		ClientSecret: clientSecret,
		Scopes:       []string{scope},
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
	}
	return s.BaseConfig, nil
}

func (s *OAuthScheme) Login(params map[string]string) (auth.Token, error) {
	conf, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	code, ok := params["code"]
	if !ok {
		return nil, ErrMissingCodeError
	}
	redirectUrl, ok := params["redirectUrl"]
	if !ok {
		return nil, ErrMissingCodeRedirectUrl
	}
	conf.RedirectURL = redirectUrl
	oauthToken, err := conf.Exchange(context.Background(), code)
	if err != nil {
		return nil, err
	}
	return s.handleToken(oauthToken)
}

func (s *OAuthScheme) handleToken(t *oauth2.Token) (*Token, error) {
	if t.AccessToken == "" {
		return nil, ErrEmptyAccessToken
	}
	conf, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	client := conf.Client(context.Background(), t)
	t0 := time.Now()
	response, err := client.Get(s.InfoUrl)
	requestLatencies.Observe(time.Since(t0).Seconds())
	if err != nil {
		requestErrors.Inc()
		return nil, err
	}
	defer response.Body.Close()
	email, err := s.Parser.Parse(response)
	if err != nil {
		return nil, err
	}
	if email == "" {
		return nil, ErrEmptyUserEmail
	}
	user, err := auth.GetUserByEmail(email)
	if err != nil {
		if err != auth.ErrUserNotFound {
			return nil, err
		}
		registrationEnabled, _ := config.GetBool("auth:user-registration")
		if !registrationEnabled {
			return nil, err
		}
		user = &auth.User{Email: email}
		err = user.Create()
		if err != nil {
			return nil, err
		}
	}
	token := Token{*t, email}
	err = token.save()
	if err != nil {
		return nil, err
	}
	return &token, nil
}

func (s *OAuthScheme) AppLogin(appName string) (auth.Token, error) {
	nativeScheme := native.NativeScheme{}
	return nativeScheme.AppLogin(appName)
}

func (s *OAuthScheme) AppLogout(token string) error {
	nativeScheme := native.NativeScheme{}
	return nativeScheme.AppLogout(token)
}

func (s *OAuthScheme) Logout(token string) error {
	return deleteToken(token)
}

func (s *OAuthScheme) Auth(header string) (auth.Token, error) {
	token, err := getToken(header)
	if err != nil {
		nativeScheme := native.NativeScheme{}
		token, nativeErr := nativeScheme.Auth(header)
		if nativeErr == nil && token.IsAppToken() {
			return token, nil
		}
		return nil, err
	}
	config, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	client := config.Client(context.Background(), &token.Token)
	t0 := time.Now()
	rsp, err := client.Get(s.InfoUrl)
	requestLatencies.Observe(time.Since(t0).Seconds())
	if err != nil {
		requestErrors.Inc()
		return nil, err
	}
	defer rsp.Body.Close()
	return token, nil
}

func (s *OAuthScheme) Name() string {
	return "oauth"
}

func (s *OAuthScheme) Info() (auth.SchemeInfo, error) {
	config, err := s.loadConfig()
	if err != nil {
		return nil, err
	}
	config.RedirectURL = "__redirect_url__"
	return auth.SchemeInfo{"authorizeUrl": config.AuthCodeURL(""), "port": strconv.Itoa(s.CallbackPort)}, nil
}

func (s *OAuthScheme) Parse(infoResponse *http.Response) (string, error) {
	user := struct {
		Email string `json:"email"`
	}{}
	data, err := ioutil.ReadAll(infoResponse.Body)
	if err != nil {
		return "", errors.Wrap(err, "unable to read user data response")
	}
	if infoResponse.StatusCode != http.StatusOK {
		return "", errors.Errorf("unexpected user data response %d: %s", infoResponse.StatusCode, data)
	}
	err = json.Unmarshal(data, &user)
	if err != nil {
		return "", errors.Wrapf(err, "unable to parse user data: %s", data)
	}
	return user.Email, nil
}

func (s *OAuthScheme) Create(user *auth.User) (*auth.User, error) {
	user.Password = ""
	err := user.Create()
	if err != nil {
		return nil, err
	}
	return user, nil
}

func (s *OAuthScheme) Remove(u *auth.User) error {
	err := deleteAllTokens(u.Email)
	if err != nil {
		return err
	}
	return u.Delete()
}
<|endoftext|>


<commit_before>package myaws

import (
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

// EC2SSHOptions customize the behavior of the SSH command.
type EC2SSHOptions struct {
	FilterTag    string
	LoginName    string
	IdentityFile string
	Private      bool
	Command      string
}

// EC2SSH resolves IP address of EC2 instance and connects to it by SSH.
func (client *Client) EC2SSH(options EC2SSHOptions) error {
	instances, err := client.FindEC2Instances(options.FilterTag, false)
	if err != nil {
		return err
	}
	if len(instances) == 0 {
		return errors.Errorf("no such instance: %s", options.FilterTag)
	}

	if len(instances) >= 2 {
		return errors.New("multiple instances found")
	}

	instance := instances[0]
	hostname, err := client.resolveEC2IPAddress(instance, options.Private)
	if err != nil {
		return errors.Wrap(err, "unable to resolve IP address:")
	}

	config, err := buildSSHConfig(options.LoginName, options.IdentityFile)
	if err != nil {
		return err
	}

	if options.Command == "" {
		return startSSHSessionWithTerminal(hostname, config)
	}

	return executeSSHCommand(hostname, config, options.Command)
}

func (client *Client) resolveEC2IPAddress(instance *ec2.Instance, private bool) (string, error) {
	if private {
		return client.resolveEC2PrivateIPAddress(instance)
	}
	return client.resolveEC2PublicIPAddress(instance)
}

func (client *Client) resolveEC2PrivateIPAddress(instance *ec2.Instance) (string, error) {
	if instance.PrivateIpAddress == nil {
		return "", errors.Errorf("no private ip address: %s", instance.InstanceId)
	}
	return *instance.PrivateIpAddress, nil
}

func (client *Client) resolveEC2PublicIPAddress(instance *ec2.Instance) (string, error) {
	if instance.PublicIpAddress == nil {
		return "", errors.Errorf("no public ip address: %s", instance.InstanceId)
	}
	return *instance.PublicIpAddress, nil
}

func buildSSHConfig(loginName string, identityFile string) (*ssh.ClientConfig, error) {
	normalizedIdentityFile := strings.Replace(identityFile, "~", os.Getenv("HOME"), 1)
	key, err := ioutil.ReadFile(normalizedIdentityFile)
	if err != nil {
		return nil, errors.Wrap(err, "unable to read private key:")
	}

	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse private key:")
	}

	config := &ssh.ClientConfig{
		User: loginName,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	}

	return config, nil
}

func buildSSHSessionPipe(session *ssh.Session) error {
	stdin, err := session.StdinPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stdin for session:")
	}
	go io.Copy(stdin, os.Stdin)

	stdout, err := session.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stdout for session:")
	}
	go io.Copy(os.Stdout, stdout)

	stderr, err := session.StderrPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stderr for session:")
	}
	go io.Copy(os.Stderr, stderr)

	return nil
}

func startSSHSessionWithTerminal(hostname string, config *ssh.ClientConfig) error {
	connection, err := ssh.Dial("tcp", hostname+":22", config)
	if err != nil {
		return errors.Wrap(err, "unable to connect:")
	}
	defer connection.Close()

	session, err := connection.NewSession()
	if err != nil {
		return errors.Wrap(err, "unable to new session failed:")
	}
	defer session.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          1,     // enable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	fd := int(os.Stdin.Fd())
	oldState, err := terminal.MakeRaw(fd)
	if err != nil {
		return errors.Wrap(err, "unable to put terminal in Raw Mode:")
	}
	defer terminal.Restore(fd, oldState)

	width, height, _ := terminal.GetSize(fd)

	if err := session.RequestPty("xterm", height, width, modes); err != nil {
		return errors.Wrap(err, "request for pseudo terminal failed:")
	}

	if err := buildSSHSessionPipe(session); err != nil {
		return err
	}

	if err := session.Shell(); err != nil {
		return errors.Wrap(err, "failed to start shell:")
	}
	session.Wait()

	return nil
}

func executeSSHCommand(hostname string, config *ssh.ClientConfig, command string) error {
	connection, err := ssh.Dial("tcp", hostname+":22", config)
	if err != nil {
		return errors.Wrap(err, "unable to connect:")
	}
	defer connection.Close()

	session, err := connection.NewSession()
	if err != nil {
		return errors.Wrap(err, "unable to new session failed:")
	}
	defer session.Close()

	if err := buildSSHSessionPipe(session); err != nil {
		return err
	}

	session.Wait()
	if err := session.Run(command); err != nil {
		return errors.Wrapf(err, "failed to execute command: %s", command)
	}

	return nil
}
<commit_msg>Fix output handling for ec2 ssh<commit_after>package myaws

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

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

// EC2SSHOptions customize the behavior of the SSH command.
type EC2SSHOptions struct {
	FilterTag    string
	LoginName    string
	IdentityFile string
	Private      bool
	Command      string
}

// EC2SSH resolves IP address of EC2 instance and connects to it by SSH.
func (client *Client) EC2SSH(options EC2SSHOptions) error {
	instances, err := client.FindEC2Instances(options.FilterTag, false)
	if err != nil {
		return err
	}
	if len(instances) == 0 {
		return errors.Errorf("no such instance: %s", options.FilterTag)
	}

	if len(instances) >= 2 {
		return errors.New("multiple instances found")
	}

	instance := instances[0]
	hostname, err := client.resolveEC2IPAddress(instance, options.Private)
	if err != nil {
		return errors.Wrap(err, "unable to resolve IP address:")
	}

	config, err := buildSSHConfig(options.LoginName, options.IdentityFile)
	if err != nil {
		return err
	}

	if options.Command == "" {
		return startSSHSessionWithTerminal(hostname, config)
	}

	return executeSSHCommand(hostname, config, options.Command)
}

func (client *Client) resolveEC2IPAddress(instance *ec2.Instance, private bool) (string, error) {
	if private {
		return client.resolveEC2PrivateIPAddress(instance)
	}
	return client.resolveEC2PublicIPAddress(instance)
}

func (client *Client) resolveEC2PrivateIPAddress(instance *ec2.Instance) (string, error) {
	if instance.PrivateIpAddress == nil {
		return "", errors.Errorf("no private ip address: %s", instance.InstanceId)
	}
	return *instance.PrivateIpAddress, nil
}

func (client *Client) resolveEC2PublicIPAddress(instance *ec2.Instance) (string, error) {
	if instance.PublicIpAddress == nil {
		return "", errors.Errorf("no public ip address: %s", instance.InstanceId)
	}
	return *instance.PublicIpAddress, nil
}

func buildSSHConfig(loginName string, identityFile string) (*ssh.ClientConfig, error) {
	normalizedIdentityFile := strings.Replace(identityFile, "~", os.Getenv("HOME"), 1)
	key, err := ioutil.ReadFile(normalizedIdentityFile)
	if err != nil {
		return nil, errors.Wrap(err, "unable to read private key:")
	}

	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse private key:")
	}

	config := &ssh.ClientConfig{
		User: loginName,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	}

	return config, nil
}

func buildSSHSessionPipe(session *ssh.Session) error {
	stdin, err := session.StdinPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stdin for session:")
	}
	go io.Copy(stdin, os.Stdin)

	stdout, err := session.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stdout for session:")
	}
	go io.Copy(os.Stdout, stdout)

	stderr, err := session.StderrPipe()
	if err != nil {
		return errors.Wrap(err, "unable to setup stderr for session:")
	}
	go io.Copy(os.Stderr, stderr)

	return nil
}

func startSSHSessionWithTerminal(hostname string, config *ssh.ClientConfig) error {
	connection, err := ssh.Dial("tcp", hostname+":22", config)
	if err != nil {
		return errors.Wrap(err, "unable to connect:")
	}
	defer connection.Close()

	session, err := connection.NewSession()
	if err != nil {
		return errors.Wrap(err, "unable to new session failed:")
	}
	defer session.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          1,     // enable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	fd := int(os.Stdin.Fd())
	oldState, err := terminal.MakeRaw(fd)
	if err != nil {
		return errors.Wrap(err, "unable to put terminal in Raw Mode:")
	}
	defer terminal.Restore(fd, oldState)

	width, height, _ := terminal.GetSize(fd)

	if err := session.RequestPty("xterm", height, width, modes); err != nil {
		return errors.Wrap(err, "request for pseudo terminal failed:")
	}

	if err := buildSSHSessionPipe(session); err != nil {
		return err
	}

	if err := session.Shell(); err != nil {
		return errors.Wrap(err, "failed to start shell:")
	}
	session.Wait()

	return nil
}

func executeSSHCommand(hostname string, config *ssh.ClientConfig, command string) error {
	connection, err := ssh.Dial("tcp", hostname+":22", config)
	if err != nil {
		return errors.Wrap(err, "unable to connect:")
	}
	defer connection.Close()

	session, err := connection.NewSession()
	if err != nil {
		return errors.Wrap(err, "unable to new session failed:")
	}
	defer session.Close()

	out, err := session.CombinedOutput(command)
	if err != nil {
		return errors.Wrapf(err, "failed to execute command: %s", command)
	}

	fmt.Println(string(out))
	return nil
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2014 The pblcache Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package spc

import (
	"github.com/pblcache/pblcache/cache"
	"github.com/pblcache/pblcache/tests"
	"os"
	"testing"
)

func TestNewSpcInfo(t *testing.T) {

	var cache *cache.Cache

	usedirectio := false
	blocksize := 4 * KB

	s := NewSpcInfo(cache, usedirectio, blocksize)
	tests.Assert(t, s.pblcache == cache)
	tests.Assert(t, s.blocksize == blocksize)
	tests.Assert(t, len(s.asus) == 3)
	tests.Assert(t, nil != s.asus[ASU1])
	tests.Assert(t, nil != s.asus[ASU2])
	tests.Assert(t, nil != s.asus[ASU3])
	tests.Assert(t, usedirectio == s.asus[ASU1].usedirectio)
	tests.Assert(t, usedirectio == s.asus[ASU2].usedirectio)
	tests.Assert(t, usedirectio == s.asus[ASU3].usedirectio)
}

func TestSpcOpen(t *testing.T) {

	// initialize
	var cache *cache.Cache
	usedirectio := false
	blocksize := 4 * KB
	s := NewSpcInfo(cache, usedirectio, blocksize)

	// Get a test file
	tmpfile := tests.Tempfile()

	// No file exists
	err := s.Open(1, tmpfile)
	tests.Assert(t, err != nil)

	// Create the file and open it
	buf := make([]byte, 16*4*KB)
	fp, err := os.Create(tmpfile)
	tests.Assert(t, err == nil)
	fp.Write(buf)
	fp.Close()
	defer os.Remove(tmpfile)

	// Now open, and it should work
	err = s.Open(1, tmpfile)
	tests.Assert(t, err == nil)
	tests.Assert(t, len(s.asus[ASU1].fps) == 1)
}
<commit_msg>pblio: TestSpcAdjustAsuSizes<commit_after>//
// Copyright (c) 2014 The pblcache Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package spc

import (
	"github.com/pblcache/pblcache/cache"
	"github.com/pblcache/pblcache/tests"
	"os"
	"testing"
)

func TestNewSpcInfo(t *testing.T) {

	var cache *cache.Cache

	usedirectio := false
	blocksize := 4 * KB

	s := NewSpcInfo(cache, usedirectio, blocksize)
	tests.Assert(t, s.pblcache == cache)
	tests.Assert(t, s.blocksize == blocksize)
	tests.Assert(t, len(s.asus) == 3)
	tests.Assert(t, nil != s.asus[ASU1])
	tests.Assert(t, nil != s.asus[ASU2])
	tests.Assert(t, nil != s.asus[ASU3])
	tests.Assert(t, usedirectio == s.asus[ASU1].usedirectio)
	tests.Assert(t, usedirectio == s.asus[ASU2].usedirectio)
	tests.Assert(t, usedirectio == s.asus[ASU3].usedirectio)
}

func TestSpcOpen(t *testing.T) {

	// initialize
	var cache *cache.Cache
	usedirectio := false
	blocksize := 4 * KB
	s := NewSpcInfo(cache, usedirectio, blocksize)

	// Get a test file
	tmpfile := tests.Tempfile()

	// No file exists
	err := s.Open(1, tmpfile)
	tests.Assert(t, err != nil)

	// Create the file and open it
	buf := make([]byte, 16*4*KB)
	fp, err := os.Create(tmpfile)
	tests.Assert(t, err == nil)
	fp.Write(buf)
	fp.Close()
	defer os.Remove(tmpfile)

	// Now open, and it should work
	err = s.Open(1, tmpfile)
	tests.Assert(t, err == nil)
}

func TestSpcAdjustAsuSizes(t *testing.T) {

	// initialize
	var cache *cache.Cache
	usedirectio := false
	blocksize := 4 * KB
	s := NewSpcInfo(cache, usedirectio, blocksize)

	// Setup some fake data
	s.asus[ASU1].len = 100
	s.asus[ASU2].len = 200
	s.asus[ASU3].len = 50
	err := s.adjustAsuSizes()

	// asu1 must be equal to asu2
	tests.Assert(t, err == nil)
	tests.Assert(t, s.asus[ASU1].len == 100)
	tests.Assert(t, s.asus[ASU2].len == 100)
	tests.Assert(t, s.asus[ASU3].len == 22)

	// Setup some fake data
	s.asus[ASU1].len = 200
	s.asus[ASU2].len = 100
	s.asus[ASU3].len = 50
	err = s.adjustAsuSizes()

	// asu1 must be equal to asu2
	tests.Assert(t, err == nil)
	tests.Assert(t, s.asus[ASU1].len == 100)
	tests.Assert(t, s.asus[ASU2].len == 100)
	tests.Assert(t, s.asus[ASU3].len == 22)

	// Setup some fake data
	s.asus[ASU1].len = 100
	s.asus[ASU2].len = 100
	s.asus[ASU3].len = 5
	err = s.adjustAsuSizes()

	// asu3 will error since it is not large enough
	tests.Assert(t, err != nil)
	tests.Assert(t, s.asus[ASU1].len == 100)
	tests.Assert(t, s.asus[ASU2].len == 100)
	tests.Assert(t, s.asus[ASU3].len == 5)

}
<|endoftext|>


<commit_before>package lz4

// #cgo CFLAGS: -O3
// #include "src/lz4hc.h"
// #include "src/lz4hc.c"
import "C"

import (
	"fmt"
)

// CompressHC compresses in and puts the content in out. len(out)
// should have enough space for the compressed data (use CompressBound
// to calculate). Returns the number of bytes in the out slice. Determines
// the compression level automatically.
func CompressHC(in, out []byte) (int, error) {
	// 0 automatically sets the compression level.
	return CompressHCLevel(in, out, 0)
}

// CompressHCLevel compresses in at the given compression level and puts the
// content in out. len(out) should have enough space for the compressed data
// (use CompressBound to calculate). Returns the number of bytes in the out
// slice. To automatically choose the compression level, use 0. Otherwise, use
// any value in the inclusive range 1 (worst) through 16 (best). Most
//applications will prefer CompressHC.
func CompressHCLevel(in, out []byte, level int) (outSize int, err error) {
	// LZ4HC does not handle empty buffers. Pass through to Compress.
	if len(in) == 0 || len(out) == 0 {
		return Compress(in, out)
	}

	outSize = int(C.LZ4_compressHC2_limitedOutput(p(in), p(out), clen(in), clen(out), C.int(level)))
	if outSize == 0 {
		err = fmt.Errorf("insufficient space for compression")
	}
	return
}
<commit_msg>comment whitespace<commit_after>package lz4

// #cgo CFLAGS: -O3
// #include "src/lz4hc.h"
// #include "src/lz4hc.c"
import "C"

import (
	"fmt"
)

// CompressHC compresses in and puts the content in out. len(out)
// should have enough space for the compressed data (use CompressBound
// to calculate). Returns the number of bytes in the out slice. Determines
// the compression level automatically.
func CompressHC(in, out []byte) (int, error) {
	// 0 automatically sets the compression level.
	return CompressHCLevel(in, out, 0)
}

// CompressHCLevel compresses in at the given compression level and puts the
// content in out. len(out) should have enough space for the compressed data
// (use CompressBound to calculate). Returns the number of bytes in the out
// slice. To automatically choose the compression level, use 0. Otherwise, use
// any value in the inclusive range 1 (worst) through 16 (best). Most
// applications will prefer CompressHC.
func CompressHCLevel(in, out []byte, level int) (outSize int, err error) {
	// LZ4HC does not handle empty buffers. Pass through to Compress.
	if len(in) == 0 || len(out) == 0 {
		return Compress(in, out)
	}

	outSize = int(C.LZ4_compressHC2_limitedOutput(p(in), p(out), clen(in), clen(out), C.int(level)))
	if outSize == 0 {
		err = fmt.Errorf("insufficient space for compression")
	}
	return
}
<|endoftext|>


<commit_before>package images

import (
    "testing"
    "fmt"
    "os"
)

func TestNegative(t *testing.T) {
    img, err := Open("./testdata/test.jpg")
    if img.Error != nil {
        t.Errorf("Error %q", err.Error())
    }
    if img.Format != "jpeg" {
        t.Errorf("Error Format, %q", img.Format)
    }
    if img.Width() != 457 || img.Height() != 343 {
        t.Errorf("Error size, %q %q", fmt.Sprintf("%v", img.Width()), fmt.Sprintf("%v", img.Height()))
    }
    if img.Quality != 100 {
        t.Errorf("Error default quality, %q", img.Quality)
    }

    img_negative := img.Negative()
    img_negative.Format = "jpeg"

    if img_negative.Error != nil {
        t.Errorf("Error negative %q", img_negative.Error.Error())
    }
    if img_negative.Format != "jpeg" {
        t.Errorf("Error negative Format, %q", img_negative.Format)
    }
    if img_negative.Width() != 457 || img_negative.Height() != 343 {
        t.Errorf("Error negative size, %q %q", fmt.Sprintf("%v", img_negative.Width()), fmt.Sprintf("%v", img_negative.Height()))
    }
    if img_negative.Quality != 100 {
        t.Errorf("Error negative default quality, %q", img_negative.Quality)
    }

    err = img_negative.Save("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative save, %q", err.Error())
    }

    img2, err := Open("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative open, %q", err.Error())
    }
    if img2.Format != "jpeg" {
        t.Errorf("Error negative Format, %q", img.Format)
    }
    if img2.Width() != 457 || img2.Height() != 343 {
        t.Errorf("Error negative open size, %q %q", fmt.Sprintf("%v", img2.Width()), fmt.Sprintf("%v", img2.Height()))
    }

    file, err := os.Open("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative open, %q", err)
    }
    hashfile := GetHashFile(file)
    if hashfile != "b4d65104a11a52df7ece664680d7db58a8ec83992b64d8f4699e7b0c2b3e1cb8" {
        t.Errorf("Error negative hash256 file, %q", hashfile)
    }
}<commit_msg>Add test for Negative func, hashfile<commit_after>package images

import (
    "testing"
    "fmt"
    "os"
)

func TestNegative(t *testing.T) {
    img, err := Open("./testdata/test.jpg")
    if img.Error != nil {
        t.Errorf("Error %q", err.Error())
    }
    if img.Format != "jpeg" {
        t.Errorf("Error Format, %q", img.Format)
    }
    if img.Width() != 457 || img.Height() != 343 {
        t.Errorf("Error size, %q %q", fmt.Sprintf("%v", img.Width()), fmt.Sprintf("%v", img.Height()))
    }
    if img.Quality != 100 {
        t.Errorf("Error default quality, %q", img.Quality)
    }

    img_negative := img.Negative()
    img_negative.Format = "jpeg"

    if img_negative.Error != nil {
        t.Errorf("Error negative %q", img_negative.Error.Error())
    }
    if img_negative.Format != "jpeg" {
        t.Errorf("Error negative Format, %q", img_negative.Format)
    }
    if img_negative.Width() != 457 || img_negative.Height() != 343 {
        t.Errorf("Error negative size, %q %q", fmt.Sprintf("%v", img_negative.Width()), fmt.Sprintf("%v", img_negative.Height()))
    }
    if img_negative.Quality != 100 {
        t.Errorf("Error negative default quality, %q", img_negative.Quality)
    }

    err = img_negative.Save("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative save, %q", err.Error())
    }

    img2, err := Open("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative open, %q", err.Error())
    }
    if img2.Format != "jpeg" {
        t.Errorf("Error negative Format, %q", img.Format)
    }
    if img2.Width() != 457 || img2.Height() != 343 {
        t.Errorf("Error negative open size, %q %q", fmt.Sprintf("%v", img2.Width()), fmt.Sprintf("%v", img2.Height()))
    }

}

func TestHashFile(t *testing.T) {
    file, err := os.Open("./testdata/test.jpg")
    if err != nil {
        t.Errorf("Error image open, %q", err)
    }
    hashfile := GetHashFile(file)
    if hashfile != "e269a4995ad439664251b38951448022706e037b40d243475f1bb3ae74329212" {
        t.Errorf("Error image hash256 file, %q", hashfile)
    }
}

func TestHashFileNegative(t *testing.T) {
    file, err := os.Open("./testdata/test_negative.jpg")
    if err != nil {
        t.Errorf("Error negative open, %q", err)
    }
    hashfile := GetHashFile(file)
    if hashfile != "b4d65104a11a52df7ece664680d7db58a8ec83992b64d8f4699e7b0c2b3e1cb8" {
        t.Errorf("Error negative hash256 file, %q", hashfile)
    }
}

<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package js

import (
	"context"
	"fmt"
	"github.com/loadimpact/k6/lib"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestSleep(t *testing.T) {
	if testing.Short() {
		return
	}

	testdata := map[string]struct {
		src string
		min time.Duration
	}{
		"float,sub-1s": {`0.2`, 200 * time.Millisecond},
		"float":        {`1.0`, 1 * time.Second},
		"int":          {`1`, 1 * time.Second},
		"exceeding":    {`5`, 2 * time.Second},
	}
	for name, data := range testdata {
		t.Run(name, func(t *testing.T) {
			r, err := newSnippetRunner(fmt.Sprintf(`
			import { sleep } from "k6";
			export default function() {
				sleep(%s);
			}`, data.src))
			assert.NoError(t, err)

			vu, err := r.NewVU()
			assert.NoError(t, err)

			ctx, _ := context.WithTimeout(context.Background(), 2*time.Second)
			start := time.Now()

			_, err = vu.RunOnce(ctx)
			assert.NoError(t, err)
			assert.True(t, time.Since(start) > data.min, "ran too short")
			assert.True(t, time.Since(start) < data.min+1*time.Second, "ran too long")
		})
	}
}

func TestDoGroup(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group } from "k6";
	export default function() {
		group("test", fn);
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	assert.NoError(t, vu.vm.Set("fn", func() {
		assert.Equal(t, "test", vu.group.Name)
	}))

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupNested(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group } from "k6";
	export default function() {
		group("outer", function() {
			group("inner", fn);
		});
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	assert.NoError(t, vu.vm.Set("fn", func() {
		assert.Equal(t, "inner", vu.group.Name)
		assert.Equal(t, "outer", vu.group.Parent.Name)
	}))

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupReturn(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, _assert } from "k6";
	export default function() {
		let v = group("group", function() {
			return 12345;
		});
		_assert(v === 12345);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupReturnTrueByDefault(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, _assert } from "k6";
	export default function() {
		let v = group("group", function() {
			// no return
		});
		_assert(v === true);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoCheck(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check } from "k6";
	export default function() {
		check(3, { "v === 3": (v) => v === 3 });
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)

	c := r.DefaultGroup.Checks["v === 3"]
	assert.NotNil(t, c)
	assert.Equal(t, "v === 3", c.Name)
	assert.Equal(t, r.DefaultGroup, c.Group)
	assert.Equal(t, int64(1), c.Passes)
	assert.Equal(t, int64(0), c.Fails)
}

func TestCheckInGroup(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, check } from "k6";
	export default function() {
		group("group", function() {
			check(3, { "v === 3": (v) => v === 3 });
		});
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)

	g := r.DefaultGroup.Groups["group"]
	assert.NotNil(t, g)
	assert.Equal(t, "group", g.Name)

	c := g.Checks["v === 3"]
	assert.NotNil(t, c)
	assert.Equal(t, "v === 3", c.Name)
	assert.Equal(t, g, c.Group)
	assert.Equal(t, int64(1), c.Passes)
	assert.Equal(t, int64(0), c.Fails)
}

func TestCheckReturnTrueOnSuccess(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check, _assert } from "k6";
	export default function() {
		let succ = check(null, { "true": true });
		_assert(succ === true);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestCheckReturnFalseAndTaintsOnFailure(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check, _assert } from "k6";
	export default function() {
		let succ = check(null, { "false": false });
		_assert(succ === false);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.Equal(t, lib.ErrVUWantsTaint, err)
}

func TestTaint(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { taint } from "k6";
	export default function() {
		taint();
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)

	_, err = vu.RunOnce(context.Background())
	assert.Equal(t, lib.ErrVUWantsTaint, err)
}
<commit_msg>[fix/test] Complaint about a possible context leak<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package js

import (
	"context"
	"fmt"
	"github.com/loadimpact/k6/lib"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestSleep(t *testing.T) {
	if testing.Short() {
		return
	}

	testdata := map[string]struct {
		src string
		min time.Duration
	}{
		"float,sub-1s": {`0.2`, 200 * time.Millisecond},
		"float":        {`1.0`, 1 * time.Second},
		"int":          {`1`, 1 * time.Second},
		"exceeding":    {`5`, 2 * time.Second},
	}
	for name, data := range testdata {
		t.Run(name, func(t *testing.T) {
			r, err := newSnippetRunner(fmt.Sprintf(`
			import { sleep } from "k6";
			export default function() {
				sleep(%s);
			}`, data.src))
			assert.NoError(t, err)

			vu, err := r.NewVU()
			assert.NoError(t, err)

			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			start := time.Now()

			_, err = vu.RunOnce(ctx)
			assert.NoError(t, err)
			assert.True(t, time.Since(start) > data.min, "ran too short")
			assert.True(t, time.Since(start) < data.min+1*time.Second, "ran too long")
		})
	}
}

func TestDoGroup(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group } from "k6";
	export default function() {
		group("test", fn);
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	assert.NoError(t, vu.vm.Set("fn", func() {
		assert.Equal(t, "test", vu.group.Name)
	}))

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupNested(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group } from "k6";
	export default function() {
		group("outer", function() {
			group("inner", fn);
		});
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	assert.NoError(t, vu.vm.Set("fn", func() {
		assert.Equal(t, "inner", vu.group.Name)
		assert.Equal(t, "outer", vu.group.Parent.Name)
	}))

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupReturn(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, _assert } from "k6";
	export default function() {
		let v = group("group", function() {
			return 12345;
		});
		_assert(v === 12345);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoGroupReturnTrueByDefault(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, _assert } from "k6";
	export default function() {
		let v = group("group", function() {
			// no return
		});
		_assert(v === true);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestDoCheck(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check } from "k6";
	export default function() {
		check(3, { "v === 3": (v) => v === 3 });
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)

	c := r.DefaultGroup.Checks["v === 3"]
	assert.NotNil(t, c)
	assert.Equal(t, "v === 3", c.Name)
	assert.Equal(t, r.DefaultGroup, c.Group)
	assert.Equal(t, int64(1), c.Passes)
	assert.Equal(t, int64(0), c.Fails)
}

func TestCheckInGroup(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { group, check } from "k6";
	export default function() {
		group("group", function() {
			check(3, { "v === 3": (v) => v === 3 });
		});
	}`)
	assert.NoError(t, err)

	vu_, err := r.NewVU()
	assert.NoError(t, err)
	vu := vu_.(*VU)

	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)

	g := r.DefaultGroup.Groups["group"]
	assert.NotNil(t, g)
	assert.Equal(t, "group", g.Name)

	c := g.Checks["v === 3"]
	assert.NotNil(t, c)
	assert.Equal(t, "v === 3", c.Name)
	assert.Equal(t, g, c.Group)
	assert.Equal(t, int64(1), c.Passes)
	assert.Equal(t, int64(0), c.Fails)
}

func TestCheckReturnTrueOnSuccess(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check, _assert } from "k6";
	export default function() {
		let succ = check(null, { "true": true });
		_assert(succ === true);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.NoError(t, err)
}

func TestCheckReturnFalseAndTaintsOnFailure(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { check, _assert } from "k6";
	export default function() {
		let succ = check(null, { "false": false });
		_assert(succ === false);
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)
	_, err = vu.RunOnce(context.Background())
	assert.Equal(t, lib.ErrVUWantsTaint, err)
}

func TestTaint(t *testing.T) {
	if testing.Short() {
		return
	}

	r, err := newSnippetRunner(`
	import { taint } from "k6";
	export default function() {
		taint();
	}`)
	assert.NoError(t, err)

	vu, err := r.NewVU()
	assert.NoError(t, err)

	_, err = vu.RunOnce(context.Background())
	assert.Equal(t, lib.ErrVUWantsTaint, err)
}
<|endoftext|>


<commit_before>package hosting

import (
	"appstax-cli/appstax/apiclient"
	"appstax-cli/appstax/fail"
	"archive/tar"
	"bufio"
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

func UploadStatic(archivePath string, progressWriter io.Writer) error {
	_, _, err := apiclient.PostFile(apiclient.Url("/appstax/hosting/static"), archivePath, progressWriter)
	return err
}

func PrepareArchive(rootPath string) (string, int64) {
	file, err := ioutil.TempFile("", "")
	fail.Handle(err)
	defer file.Close()
	fileWriter := bufio.NewWriter(file)
	defer fileWriter.Flush()
	gzipWriter, err := gzip.NewWriterLevel(fileWriter, gzip.BestCompression)
	fail.Handle(err)
	defer gzipWriter.Close()
	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	fullRootPath, err := filepath.Abs(rootPath)
	fail.Handle(err)
	addAllToArchive(fullRootPath, tarWriter)

	tarWriter.Close()
	gzipWriter.Close()
	fileWriter.Flush()
	file.Close()

	fileInfo, err := os.Stat(file.Name())
	fail.Handle(err)
	return file.Name(), fileInfo.Size()
}

func addAllToArchive(fullRootPath string, tarWriter *tar.Writer) {
	filepath.Walk(fullRootPath, func(path string, fileInfo os.FileInfo, err error) error {
		fail.Handle(err)
		if !fileInfo.IsDir() && fileInfo.Name()[:1] != "." {
			addFileToArchive(path, path[len(fullRootPath+"/"):], tarWriter, fileInfo)
		}
		return nil
	})
}

func addFileToArchive(filePath string, addPath string, tarWriter *tar.Writer, fileInfo os.FileInfo) {
	//println("Adding " + addPath + " (" + filePath + ")")
	fileReader, err := os.Open(filePath)
	fail.Handle(err)
	defer fileReader.Close()

	header := new(tar.Header)
	header.Name = addPath
	header.Size = fileInfo.Size()
	header.Mode = int64(fileInfo.Mode())
	header.ModTime = fileInfo.ModTime()

	err = tarWriter.WriteHeader(header)
	fail.Handle(err)
	_, err = io.Copy(tarWriter, fileReader)
	fail.Handle(err)
}
<commit_msg>Add debug logging to hosting.PrepareArchive()<commit_after>package hosting

import (
	"appstax-cli/appstax/apiclient"
	"appstax-cli/appstax/fail"
	"appstax-cli/appstax/log"
	"archive/tar"
	"bufio"
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

func UploadStatic(archivePath string, progressWriter io.Writer) error {
	_, _, err := apiclient.PostFile(apiclient.Url("/appstax/hosting/static"), archivePath, progressWriter)
	return err
}

func PrepareArchive(rootPath string) (string, int64) {
	file, err := ioutil.TempFile("", "")
	fail.Handle(err)
	defer file.Close()
	fileWriter := bufio.NewWriter(file)
	defer fileWriter.Flush()
	gzipWriter, err := gzip.NewWriterLevel(fileWriter, gzip.BestCompression)
	fail.Handle(err)
	defer gzipWriter.Close()
	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	fullRootPath, err := filepath.Abs(rootPath)
	fail.Handle(err)
	addAllToArchive(fullRootPath, tarWriter)

	tarWriter.Close()
	gzipWriter.Close()
	fileWriter.Flush()
	file.Close()

	fileInfo, err := os.Stat(file.Name())
	fail.Handle(err)
	return file.Name(), fileInfo.Size()
}

func addAllToArchive(fullRootPath string, tarWriter *tar.Writer) {
	log.Debugf("Creating archive by walking from root path %s", fullRootPath)
	filepath.Walk(fullRootPath, func(path string, fileInfo os.FileInfo, err error) error {
		fail.Handle(err)
		if !fileInfo.IsDir() && fileInfo.Name()[:1] != "." {
			addFileToArchive(path, path[len(fullRootPath+"/"):], tarWriter, fileInfo)
		} else {
			log.Debugf("Ignoring path %s", path)
		}
		return nil
	})
}

func addFileToArchive(filePath string, addPath string, tarWriter *tar.Writer, fileInfo os.FileInfo) {
	log.Debugf("Adding file %s", addPath)
	fileReader, err := os.Open(filePath)
	fail.Handle(err)
	defer fileReader.Close()

	header := new(tar.Header)
	header.Name = addPath
	header.Size = fileInfo.Size()
	header.Mode = int64(fileInfo.Mode())
	header.ModTime = fileInfo.ModTime()

	err = tarWriter.WriteHeader(header)
	fail.Handle(err)
	_, err = io.Copy(tarWriter, fileReader)
	fail.Handle(err)
}
<|endoftext|>


<commit_before>package gomail

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/mail"
	"net/smtp"
	"strings"
)

// A Mailer represents an SMTP server.
type Mailer struct {
	addr   string
	host   string
	config *tls.Config
	auth   smtp.Auth
	send   SendMailFunc
}

// A MailerSetting can be used in a mailer constructor to configure it.
type MailerSetting func(m *Mailer)

// SetSendMail allows to set the email-sending function of a mailer.
//
// Example:
//
//	myFunc := func(addr string, a smtp.Auth, from string, to []string, msg []byte) error {
//		// Implement your email-sending function similar to smtp.SendMail
//	}
//	mailer := gomail.NewMailer("host", "user", "pwd", 465, SetSendMail(myFunc))
func SetSendMail(s SendMailFunc) MailerSetting {
	return func(m *Mailer) {
		m.send = s
	}
}

// SetTLSConfig allows to set the TLS configuration used to connect the SMTP
// server.
func SetTLSConfig(c *tls.Config) MailerSetting {
	return func(m *Mailer) {
		m.config = c
	}
}

// A SendMailFunc is a function to send emails with the same signature than
// smtp.SendMail.
type SendMailFunc func(addr string, a smtp.Auth, from string, to []string, msg []byte) error

// NewMailer returns a mailer. The given parameters are used to connect to the
// SMTP server via a PLAIN authentication mechanism.
func NewMailer(host string, username string, password string, port int, settings ...MailerSetting) *Mailer {
	return NewCustomMailer(
		fmt.Sprintf("%s:%d", host, port),
		smtp.PlainAuth("", username, password, host),
		settings...,
	)
}

// NewCustomMailer creates a mailer with the given authentication mechanism.
//
// Example:
//
//	gomail.NewCustomMailer("host:587", smtp.CRAMMD5Auth("username", "secret"))
func NewCustomMailer(addr string, auth smtp.Auth, settings ...MailerSetting) *Mailer {
	// Error is not handled here to preserve backward compatibility
	host, port, _ := net.SplitHostPort(addr)

	m := &Mailer{
		addr: addr,
		host: host,
		auth: auth,
	}

	for _, s := range settings {
		s(m)
	}

	if m.config == nil {
		m.config = &tls.Config{ServerName: host}
	}
	if m.send == nil {
		m.send = m.getSendMailFunc(port == "465")
	}

	return m
}

// Send sends the emails to all the recipients of the message.
func (m *Mailer) Send(msg *Message) error {
	message := msg.Export()

	from, err := getFrom(message)
	if err != nil {
		return err
	}
	recipients, bcc, err := getRecipients(message)
	if err != nil {
		return err
	}

	h := flattenHeader(message, "")
	body, err := ioutil.ReadAll(message.Body)
	if err != nil {
		return err
	}

	mail := append(h, body...)
	if err := m.send(m.addr, m.auth, from, recipients, mail); err != nil {
		return err
	}

	for _, to := range bcc {
		h = flattenHeader(message, to)
		mail = append(h, body...)
		if err := m.send(m.addr, m.auth, from, []string{to}, mail); err != nil {
			return err
		}
	}

	return nil
}

func flattenHeader(msg *mail.Message, bcc string) []byte {
	buf := getBuffer()
	defer putBuffer(buf)

	for field, value := range msg.Header {
		if field != "Bcc" {
			buf.WriteString(field)
			buf.WriteString(": ")
			buf.WriteString(strings.Join(value, ", "))
			buf.WriteString("\r\n")
		} else if bcc != "" {
			for _, to := range value {
				if strings.Contains(to, bcc) {
					buf.WriteString(field)
					buf.WriteString(": ")
					buf.WriteString(to)
					buf.WriteString("\r\n")
				}
			}
		}
	}
	buf.WriteString("\r\n")

	return buf.Bytes()
}

func getFrom(msg *mail.Message) (string, error) {
	from := msg.Header.Get("Sender")
	if from == "" {
		from = msg.Header.Get("From")
		if from == "" {
			return "", errors.New("mailer: invalid message, \"From\" field is absent")
		}
	}

	return parseAddress(from)
}

func getRecipients(msg *mail.Message) (recipients, bcc []string, err error) {
	for _, field := range []string{"Bcc", "To", "Cc"} {
		if addresses, ok := msg.Header[field]; ok {
			for _, addr := range addresses {
				switch field {
				case "Bcc":
					bcc, err = addAdress(bcc, addr)
				default:
					recipients, err = addAdress(recipients, addr)
				}
				if err != nil {
					return recipients, bcc, err
				}
			}
		}
	}

	return recipients, bcc, nil
}

func addAdress(list []string, addr string) ([]string, error) {
	addr, err := parseAddress(addr)
	if err != nil {
		return list, err
	}
	for _, a := range list {
		if addr == a {
			return list, nil
		}
	}

	return append(list, addr), nil
}

func parseAddress(field string) (string, error) {
	a, err := mail.ParseAddress(field)
	if a == nil {
		return "", err
	}

	return a.Address, err
}
<commit_msg>Bugfix. BCC should not be added in header<commit_after>package gomail

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/mail"
	"net/smtp"
	"strings"
)

// A Mailer represents an SMTP server.
type Mailer struct {
	addr   string
	host   string
	config *tls.Config
	auth   smtp.Auth
	send   SendMailFunc
}

// A MailerSetting can be used in a mailer constructor to configure it.
type MailerSetting func(m *Mailer)

// SetSendMail allows to set the email-sending function of a mailer.
//
// Example:
//
//	myFunc := func(addr string, a smtp.Auth, from string, to []string, msg []byte) error {
//		// Implement your email-sending function similar to smtp.SendMail
//	}
//	mailer := gomail.NewMailer("host", "user", "pwd", 465, SetSendMail(myFunc))
func SetSendMail(s SendMailFunc) MailerSetting {
	return func(m *Mailer) {
		m.send = s
	}
}

// SetTLSConfig allows to set the TLS configuration used to connect the SMTP
// server.
func SetTLSConfig(c *tls.Config) MailerSetting {
	return func(m *Mailer) {
		m.config = c
	}
}

// A SendMailFunc is a function to send emails with the same signature than
// smtp.SendMail.
type SendMailFunc func(addr string, a smtp.Auth, from string, to []string, msg []byte) error

// NewMailer returns a mailer. The given parameters are used to connect to the
// SMTP server via a PLAIN authentication mechanism.
func NewMailer(host string, username string, password string, port int, settings ...MailerSetting) *Mailer {
	return NewCustomMailer(
		fmt.Sprintf("%s:%d", host, port),
		smtp.PlainAuth("", username, password, host),
		settings...,
	)
}

// NewCustomMailer creates a mailer with the given authentication mechanism.
//
// Example:
//
//	gomail.NewCustomMailer("host:587", smtp.CRAMMD5Auth("username", "secret"))
func NewCustomMailer(addr string, auth smtp.Auth, settings ...MailerSetting) *Mailer {
	// Error is not handled here to preserve backward compatibility
	host, port, _ := net.SplitHostPort(addr)

	m := &Mailer{
		addr: addr,
		host: host,
		auth: auth,
	}

	for _, s := range settings {
		s(m)
	}

	if m.config == nil {
		m.config = &tls.Config{ServerName: host}
	}
	if m.send == nil {
		m.send = m.getSendMailFunc(port == "465")
	}

	return m
}

// Send sends the emails to all the recipients of the message.
func (m *Mailer) Send(msg *Message) error {
	message := msg.Export()

	from, err := getFrom(message)
	if err != nil {
		return err
	}
	recipients, bcc, err := getRecipients(message)
	if err != nil {
		return err
	}

	h := flattenHeader(message, "")
	body, err := ioutil.ReadAll(message.Body)
	if err != nil {
		return err
	}

	mail := append(h, body...)
	if err := m.send(m.addr, m.auth, from, recipients, mail); err != nil {
		return err
	}

	for _, to := range bcc {
		h = flattenHeader(message, to)
		mail = append(h, body...)
		if err := m.send(m.addr, m.auth, from, []string{to}, mail); err != nil {
			return err
		}
	}

	return nil
}

func flattenHeader(msg *mail.Message, bcc string) []byte {
	buf := getBuffer()
	defer putBuffer(buf)

	for field, value := range msg.Header {
		if field != "Bcc" {
			buf.WriteString(field)
			buf.WriteString(": ")
			buf.WriteString(strings.Join(value, ", "))
			buf.WriteString("\r\n")
		}
	}
	buf.WriteString("\r\n")

	return buf.Bytes()
}

func getFrom(msg *mail.Message) (string, error) {
	from := msg.Header.Get("Sender")
	if from == "" {
		from = msg.Header.Get("From")
		if from == "" {
			return "", errors.New("mailer: invalid message, \"From\" field is absent")
		}
	}

	return parseAddress(from)
}

func getRecipients(msg *mail.Message) (recipients, bcc []string, err error) {
	for _, field := range []string{"Bcc", "To", "Cc"} {
		if addresses, ok := msg.Header[field]; ok {
			for _, addr := range addresses {
				switch field {
				case "Bcc":
					bcc, err = addAdress(bcc, addr)
				default:
					recipients, err = addAdress(recipients, addr)
				}
				if err != nil {
					return recipients, bcc, err
				}
			}
		}
	}

	return recipients, bcc, nil
}

func addAdress(list []string, addr string) ([]string, error) {
	addr, err := parseAddress(addr)
	if err != nil {
		return list, err
	}
	for _, a := range list {
		if addr == a {
			return list, nil
		}
	}

	return append(list, addr), nil
}

func parseAddress(field string) (string, error) {
	a, err := mail.ParseAddress(field)
	if a == nil {
		return "", err
	}

	return a.Address, err
}
<|endoftext|>


<commit_before>package database

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/url"
	"time"

	lru "github.com/hashicorp/golang-lru"
	_ "github.com/lib/pq"

	"github.com/mozilla/tls-observatory/connection"
)

type DB struct {
	*sql.DB
	paths *lru.ARCCache
}

type Scan struct {
	ID               int64                  `json:"id"`
	Timestamp        time.Time              `json:"timestamp"`
	Target           string                 `json:"target"`
	Replay           int                    `json:"replay"` //hours or days
	Has_tls          bool                   `json:"has_tls"`
	Cert_id          int64                  `json:"cert_id"`
	Trust_id         int64                  `json:"trust_id"`
	Is_valid         bool                   `json:"is_valid"`
	Validation_error string                 `json:"validation_error,omitempty"`
	Complperc        int                    `json:"completion_perc"`
	Conn_info        connection.Stored      `json:"connection_info"`
	AnalysisResults  Analyses               `json:"analysis,omitempty"`
	Ack              bool                   `json:"ack"`
	Attempts         int                    `json:"attempts"` //number of retries
	AnalysisParams   map[string]interface{} `json:"analysis_params"`
}

type Analysis struct {
	ID       int64           `json:"id"`
	Analyzer string          `json:"analyzer"`
	Result   json.RawMessage `json:"result"`
	Success  bool            `json:"success"`
}

type Analyses []Analysis

// Methods used for the sorting of analyses

func (slice Analyses) Len() int {
	return len(slice)
}

func (slice Analyses) Less(i, j int) bool {
	return slice[i].Analyzer < slice[j].Analyzer
}

func (slice Analyses) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func RegisterConnection(dbname, user, password, hostport, sslmode string) (*DB, error) {
	userPass := url.UserPassword(user, password)
	url := fmt.Sprintf("postgres://%s@%s/%s?sslmode=%s",
		userPass.String(), hostport, dbname, sslmode)
	dbfd, err := sql.Open("postgres", url)
	if err != nil {
		return nil, err
	}
	paths, err := lru.NewARC(10000)
	return &DB{dbfd, paths}, nil
}

func (db *DB) NewScan(domain string, rplay int, jsonParams []byte) (Scan, error) {
	timestamp := time.Now().UTC()

	var id int64

	err := db.QueryRow(`INSERT INTO scans
			(timestamp, target, replay, has_tls, is_valid, completion_perc, validation_error, conn_info, ack, attempts, analysis_params)
			VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
			RETURNING id`,
		timestamp, domain, rplay, false, false, 0, "", []byte("null"), false, 1, jsonParams).Scan(&id)

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

	return Scan{ID: id, Timestamp: timestamp, Replay: rplay}, nil
}

// GetLastScanTimeForTarget searches the database for the latest scan for a specific target.
// It returns both the scan timestamp and the id of the scan to enable the api to
// respond to clients with just one db query.
func (db *DB) GetLastScanTimeForTarget(target string) (int64, time.Time, error) {

	var (
		id int64
		t  time.Time
	)

	row := db.QueryRow(`SELECT id, timestamp
			    FROM scans
			    WHERE target=$1 ORDER BY timestamp DESC LIMIT 1`, target)

	err := row.Scan(&id, &t)

	if err != nil {
		if err == sql.ErrNoRows {
			return -1, time.Now(), nil
		} else {
			return -1, time.Now(), err
		}
	}
	return id, t, nil
}

func (db *DB) GetScanByID(id int64) (Scan, error) {

	s := Scan{}
	s.ID = id

	var cID, tID sql.NullInt64

	var isvalid sql.NullBool

	var ci []byte
	var params []byte

	row := db.QueryRow(`SELECT timestamp, target, replay, has_tls, cert_id, trust_id,
				   is_valid, completion_perc, validation_error, conn_info, ack, attempts, analysis_params
			    FROM scans
			    WHERE id=$1`, id)

	err := row.Scan(&s.Timestamp, &s.Target, &s.Replay, &s.Has_tls, &cID, &tID,
		&isvalid, &s.Complperc, &s.Validation_error, &ci, &s.Ack, &s.Attempts, &params)

	if err != nil {
		if err == sql.ErrNoRows {
			s.ID = -1
			return s, nil
		} else {
			return s, err
		}
	}

	if cID.Valid {
		s.Cert_id = cID.Int64

	} else {
		s.Cert_id = -1
	}

	if tID.Valid {
		s.Trust_id = tID.Int64
	} else {
		s.Trust_id = -1
	}

	if isvalid.Valid {
		s.Is_valid = isvalid.Bool
	} else {
		s.Is_valid = false
	}

	err = json.Unmarshal(ci, &s.Conn_info)
	if err != nil {
		return s, err
	}

	err = json.Unmarshal(params, &s.AnalysisParams)
	if err != nil {
		return s, err
	}

	if s.Complperc > 40 {
		s.AnalysisResults, err = db.GetAnalysisByScan(s.ID)
		return s, err
	}

	return s, nil
}

func (db *DB) GetAnalysisByScan(id int64) ([]Analysis, error) {

	var ana []Analysis
	rows, err := db.Query("SELECT id,worker_name,output,success FROM analysis WHERE scan_id=$1", id)
	if err != nil {
		return ana, err
	}
	defer rows.Close()
	for rows.Next() {
		a := Analysis{}
		if err := rows.Scan(&a.ID, &a.Analyzer, &a.Result, &a.Success); err != nil {
			return ana, err
		}
		ana = append(ana, a)
	}
	return ana, nil
}

func (db *DB) UpdateScanCompletionPercentage(id int64, p int) error {
	_, err := db.Exec("UPDATE scans SET completion_perc=$1 WHERE id=$2", p, id)

	return err
}

func (db *DB) InsertWorkerAnalysis(scanid int64, jsonRes []byte, workerName string) error {
	_, err := db.Exec("INSERT INTO analysis(scan_id,worker_name,output) VALUES($1,$2,$3)", scanid, workerName, jsonRes)

	return err
}
<commit_msg>Make scan attempts start at 0<commit_after>package database

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/url"
	"time"

	lru "github.com/hashicorp/golang-lru"
	_ "github.com/lib/pq"

	"github.com/mozilla/tls-observatory/connection"
)

type DB struct {
	*sql.DB
	paths *lru.ARCCache
}

type Scan struct {
	ID               int64                  `json:"id"`
	Timestamp        time.Time              `json:"timestamp"`
	Target           string                 `json:"target"`
	Replay           int                    `json:"replay"` //hours or days
	Has_tls          bool                   `json:"has_tls"`
	Cert_id          int64                  `json:"cert_id"`
	Trust_id         int64                  `json:"trust_id"`
	Is_valid         bool                   `json:"is_valid"`
	Validation_error string                 `json:"validation_error,omitempty"`
	Complperc        int                    `json:"completion_perc"`
	Conn_info        connection.Stored      `json:"connection_info"`
	AnalysisResults  Analyses               `json:"analysis,omitempty"`
	Ack              bool                   `json:"ack"`
	Attempts         int                    `json:"attempts"` //number of retries
	AnalysisParams   map[string]interface{} `json:"analysis_params"`
}

type Analysis struct {
	ID       int64           `json:"id"`
	Analyzer string          `json:"analyzer"`
	Result   json.RawMessage `json:"result"`
	Success  bool            `json:"success"`
}

type Analyses []Analysis

// Methods used for the sorting of analyses

func (slice Analyses) Len() int {
	return len(slice)
}

func (slice Analyses) Less(i, j int) bool {
	return slice[i].Analyzer < slice[j].Analyzer
}

func (slice Analyses) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func RegisterConnection(dbname, user, password, hostport, sslmode string) (*DB, error) {
	userPass := url.UserPassword(user, password)
	url := fmt.Sprintf("postgres://%s@%s/%s?sslmode=%s",
		userPass.String(), hostport, dbname, sslmode)
	dbfd, err := sql.Open("postgres", url)
	if err != nil {
		return nil, err
	}
	paths, err := lru.NewARC(10000)
	return &DB{dbfd, paths}, nil
}

func (db *DB) NewScan(domain string, rplay int, jsonParams []byte) (Scan, error) {
	timestamp := time.Now().UTC()

	var id int64

	err := db.QueryRow(`INSERT INTO scans
			(timestamp, target, replay, has_tls, is_valid, completion_perc, validation_error, conn_info, ack, attempts, analysis_params)
			VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
			RETURNING id`,
		timestamp, domain, rplay, false, false, 0, "", []byte("null"), false, 0, jsonParams).Scan(&id)

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

	return Scan{ID: id, Timestamp: timestamp, Replay: rplay}, nil
}

// GetLastScanTimeForTarget searches the database for the latest scan for a specific target.
// It returns both the scan timestamp and the id of the scan to enable the api to
// respond to clients with just one db query.
func (db *DB) GetLastScanTimeForTarget(target string) (int64, time.Time, error) {

	var (
		id int64
		t  time.Time
	)

	row := db.QueryRow(`SELECT id, timestamp
			    FROM scans
			    WHERE target=$1 ORDER BY timestamp DESC LIMIT 1`, target)

	err := row.Scan(&id, &t)

	if err != nil {
		if err == sql.ErrNoRows {
			return -1, time.Now(), nil
		} else {
			return -1, time.Now(), err
		}
	}
	return id, t, nil
}

func (db *DB) GetScanByID(id int64) (Scan, error) {

	s := Scan{}
	s.ID = id

	var cID, tID sql.NullInt64

	var isvalid sql.NullBool

	var ci []byte
	var params []byte

	row := db.QueryRow(`SELECT timestamp, target, replay, has_tls, cert_id, trust_id,
				   is_valid, completion_perc, validation_error, conn_info, ack, attempts, analysis_params
			    FROM scans
			    WHERE id=$1`, id)

	err := row.Scan(&s.Timestamp, &s.Target, &s.Replay, &s.Has_tls, &cID, &tID,
		&isvalid, &s.Complperc, &s.Validation_error, &ci, &s.Ack, &s.Attempts, &params)

	if err != nil {
		if err == sql.ErrNoRows {
			s.ID = -1
			return s, nil
		} else {
			return s, err
		}
	}

	if cID.Valid {
		s.Cert_id = cID.Int64

	} else {
		s.Cert_id = -1
	}

	if tID.Valid {
		s.Trust_id = tID.Int64
	} else {
		s.Trust_id = -1
	}

	if isvalid.Valid {
		s.Is_valid = isvalid.Bool
	} else {
		s.Is_valid = false
	}

	err = json.Unmarshal(ci, &s.Conn_info)
	if err != nil {
		return s, err
	}

	err = json.Unmarshal(params, &s.AnalysisParams)
	if err != nil {
		return s, err
	}

	if s.Complperc > 40 {
		s.AnalysisResults, err = db.GetAnalysisByScan(s.ID)
		return s, err
	}

	return s, nil
}

func (db *DB) GetAnalysisByScan(id int64) ([]Analysis, error) {

	var ana []Analysis
	rows, err := db.Query("SELECT id,worker_name,output,success FROM analysis WHERE scan_id=$1", id)
	if err != nil {
		return ana, err
	}
	defer rows.Close()
	for rows.Next() {
		a := Analysis{}
		if err := rows.Scan(&a.ID, &a.Analyzer, &a.Result, &a.Success); err != nil {
			return ana, err
		}
		ana = append(ana, a)
	}
	return ana, nil
}

func (db *DB) UpdateScanCompletionPercentage(id int64, p int) error {
	_, err := db.Exec("UPDATE scans SET completion_perc=$1 WHERE id=$2", p, id)

	return err
}

func (db *DB) InsertWorkerAnalysis(scanid int64, jsonRes []byte, workerName string) error {
	_, err := db.Exec("INSERT INTO analysis(scan_id,worker_name,output) VALUES($1,$2,$3)", scanid, workerName, jsonRes)

	return err
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/kataras/iris"
)

func main() {
	iris.Config.Render.Template.Layout = "layouts/layout.html" // default ""
	//iris.Config.Render.Template.Gzip = true
	iris.Get("/", func(ctx *iris.Context) {
		if err := ctx.Render("page1.html", nil); err != nil {
			println(err.Error())
		}
	})

	// remove the layout for a specific route
	iris.Get("/nolayout", func(ctx *iris.Context) {
		if err := ctx.Render("page1.html", nil, iris.NoLayout); err != nil {
			println(err.Error())
		}
	})

	// set or remove a layout for a party
	my := iris.Party("/my").Layout("layouts/mylayout.html")
	{
		my.Get("/", func(ctx *iris.Context) {
			ctx.MustRender("page1.html", nil)
		})
		my.Get("/other", func(ctx *iris.Context) {
			ctx.MustRender("page1.html", nil)
		})
	}

	iris.Listen(":8080")
}
<commit_msg>add a note for party layout<commit_after>package main

import (
	"github.com/kataras/iris"
)

func main() {
	iris.Config.Render.Template.Layout = "layouts/layout.html" // default ""
	//iris.Config.Render.Template.Gzip = true
	iris.Get("/", func(ctx *iris.Context) {
		if err := ctx.Render("page1.html", nil); err != nil {
			println(err.Error())
		}
	})

	// remove the layout for a specific route
	iris.Get("/nolayout", func(ctx *iris.Context) {
		if err := ctx.Render("page1.html", nil, iris.NoLayout); err != nil {
			println(err.Error())
		}
	})

	// set a layout for a party, .Layout should be BEFORE any Get or other Handle party's method
	my := iris.Party("/my").Layout("layouts/mylayout.html")
	{
		my.Get("/", func(ctx *iris.Context) {
			ctx.MustRender("page1.html", nil)
		})
		my.Get("/other", func(ctx *iris.Context) {
			ctx.MustRender("page1.html", nil)
		})
	}

	iris.Listen(":8080")
}
<|endoftext|>


<commit_before>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.49 Safari/537.36"

	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
	subRE          = dns.AnySubdomainRegex()
	crawlRE        = regexp.MustCompile(`\.\w{3,4}($|\?)`)
	crawlFileTypes = []string{".html", ".htm", "xhtml", ".js", ".php"}
	nameStripRE    = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			for _, n := range subRE.FindAllString(string(r.Body), -1) {
				if name := CleanName(n); whichDomain(name, scope) != "" {
					m.Lock()
					results.Insert(name)
					m.Unlock()
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.ToLower(ext)

		var found bool
		for _, t := range crawlFileTypes {
			if ext == t {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<commit_msg>Enhanced web crawling feature<commit_after>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.49 Safari/537.36"

	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
	subRE           = dns.AnySubdomainRegex()
	crawlFileEnds   = []string{"html", "do", "action", "cgi"}
	crawlFileStarts = []string{"js", "htm", "as", "php", "inc"}
	crawlRE         = regexp.MustCompile(`\.[a-z0-9]{2,6}($|\?|#)`)
	nameStripRE     = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			resp, err := httputil.DumpResponse(interface{}(r).(*http.Response), true)
			if err == nil {
				for _, n := range subRE.FindAllString(string(resp), -1) {
					if name := CleanName(n); whichDomain(name, scope) != "" {
						m.Lock()
						results.Insert(name)
						m.Unlock()
					}
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.TrimRight(ext, "?#")

		var found bool
		for _, s := range crawlFileStarts {
			if strings.HasPrefix(ext, "." + s) {
				found = true
				break
			}
		}
		for _, e := range crawlFileEnds {
			if strings.HasSuffix(ext, e) {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<|endoftext|>