<commit_before>// Copyright 2020 Asim Aslam
//
// 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
//
//     https://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.
//
// Original source: github.com/micro/go-micro/v3/auth/auth.go

package auth

import (
	"context"
	"errors"
	"time"
)

var (
	// DefaultAuth implementation
	DefaultAuth Auth
	// ErrInvalidToken is when the token provided is not valid
	ErrInvalidToken = errors.New("invalid token provided")
	// ErrForbidden is when a user does not have the necessary scope to access a resource
	ErrForbidden = errors.New("resource forbidden")
)

const (
	// ScopePublic is the scope applied to a rule to allow access to the public
	ScopePublic = ""
	// ScopeAccount is the scope applied to a rule to limit to users with any valid account
	ScopeAccount = "*"
	// ScopeAnyNamespaceAccount is the scope applied to a rule to limit to users with any valid account regardless of issuer
	ScopeAnyNamespaceAccount = "+"
)

// Account provided by an auth provider
type Account struct {
	// ID of the account e.g. UUID. Should not change
	ID string `json:"id"`
	// Type of the account, e.g. service
	Type string `json:"type"`
	// Issuer of the account
	Issuer string `json:"issuer"`
	// Any other associated metadata
	Metadata map[string]string `json:"metadata"`
	// Scopes the account has access to
	Scopes []string `json:"scopes"`
	// Secret for the account, e.g. the password
	Secret string `json:"secret"`
	// Name of the account. User friendly name that might change e.g. a username or email
	Name string `json:"name"`
}

// AccountToken can be short or long lived
type AccountToken struct {
	// The token to be used for accessing resources
	AccessToken string `json:"access_token"`
	// RefreshToken to be used to generate a new token
	RefreshToken string `json:"refresh_token"`
	// Time of token creation
	Created time.Time `json:"created"`
	// Time of token expiry
	Expiry time.Time `json:"expiry"`
}

// Expired returns a boolean indicating if the token needs to be refreshed
func (t *AccountToken) Expired() bool {
	return t.Expiry.Unix() < time.Now().Unix()
}

// Resource is an entity such as a user or
type Resource struct {
	// Name of the resource, e.g. go.micro.service.notes
	Name string `json:"name"`
	// Type of resource, e.g. service
	Type string `json:"type"`
	// Endpoint resource e.g NotesService.Create
	Endpoint string `json:"endpoint"`
}

// Access defines the type of access a rule grants
type Access int

const (
	// AccessGranted to a resource
	AccessGranted Access = iota
	// AccessDenied to a resource
	AccessDenied
)

// Rule is used to verify access to a resource
type Rule struct {
	// ID of the rule, e.g. "public"
	ID string
	// Scope the rule requires, a blank scope indicates open to the public and * indicates the rule
	// applies to any valid account
	Scope string
	// Resource the rule applies to
	Resource *Resource
	// Access determines if the rule grants or denies access to the resource
	Access Access
	// Priority the rule should take when verifying a request, the higher the value the sooner the
	// rule will be applied
	Priority int32
}

// Auth provides authentication and authorization
type Auth interface {
	// Init the auth
	Init(opts ...Option)
	// Options set for auth
	Options() Options
	// Generate a new account
	Generate(id string, opts ...GenerateOption) (*Account, error)
	// Verify an account has access to a resource using the rules
	Verify(acc *Account, res *Resource, opts ...VerifyOption) error
	// Inspect a token
	Inspect(token string) (*Account, error)
	// Token generated using refresh token or credentials
	Token(opts ...TokenOption) (*AccountToken, error)
	// Grant access to a resource
	Grant(rule *Rule) error
	// Revoke access to a resource
	Revoke(rule *Rule) error
	// Rules returns all the rules used to verify requests
	Rules(...RulesOption) ([]*Rule, error)
	// String returns the name of the implementation
	String() string
}

// Generate a new account
func Generate(id string, opts ...GenerateOption) (*Account, error) {
	return DefaultAuth.Generate(id, opts...)
}

// Verify an account has access to a resource using the rules
func Verify(acc *Account, res *Resource, opts ...VerifyOption) error {
	return DefaultAuth.Verify(acc, res, opts...)
}

// Inspect a token
func Inspect(token string) (*Account, error) {
	return DefaultAuth.Inspect(token)
}

// Token generated using refresh token or credentials
func Token(opts ...TokenOption) (*AccountToken, error) {
	return DefaultAuth.Token(opts...)
}

// Grant access to a resource
func Grant(rule *Rule) error {
	return DefaultAuth.Grant(rule)
}

// Revoke access to a resource
func Revoke(rule *Rule) error {
	return DefaultAuth.Revoke(rule)
}

// Rules returns all the rules used to verify requests
func Rules(...RulesOption) ([]*Rule, error) {
	return DefaultAuth.Rules()
}

type accountKey struct{}

// AccountFromContext gets the account from the context, which
// is set by the auth wrapper at the start of a call. If the account
// is not set, a nil account will be returned. The error is only returned
// when there was a problem retrieving an account
func AccountFromContext(ctx context.Context) (*Account, bool) {
	acc, ok := ctx.Value(accountKey{}).(*Account)
	return acc, ok
}

// ContextWithAccount sets the account in the context
func ContextWithAccount(ctx context.Context, account *Account) context.Context {
	return context.WithValue(ctx, accountKey{}, account)
}
<commit_msg>Update auth.go<commit_after>// Copyright 2020 Asim Aslam
//
// 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
//
//     https://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.
//
// Original source: github.com/micro/go-micro/v3/auth/auth.go

package auth

import (
	"context"
	"errors"
	"time"
)

var (
	// DefaultAuth implementation
	DefaultAuth Auth
	// ErrInvalidToken is when the token provided is not valid
	ErrInvalidToken = errors.New("invalid token provided")
	// ErrForbidden is when a user does not have the necessary scope to access a resource
	ErrForbidden = errors.New("resource forbidden")
)

const (
	// ScopePublic is the scope applied to a rule to allow access to the public
	ScopePublic = ""
	// ScopeAccount is the scope applied to a rule to limit to users with any valid account
	ScopeAccount = "*"
	// ScopeAnyNamespaceAccount is the scope applied to a rule to limit to users with any valid account regardless of issuer
	ScopeAnyNamespaceAccount = "+"
)

// Account provided by an auth provider
type Account struct {
	// ID of the account e.g. UUID. Should not change
	ID string `json:"id"`
	// Type of the account, e.g. cient, service, user
	Type string `json:"type"`
	// Issuer of the account e.g micro
	Issuer string `json:"issuer"`
	// Any other associated metadata
	Metadata map[string]string `json:"metadata"`
	// Scopes the account has access to
	Scopes []string `json:"scopes"`
	// Secret for the account, e.g. the password
	Secret string `json:"secret"`
	// Name of the account. User friendly name that might change e.g. a username or email
	Name string `json:"name"`
}

// AccountToken can be short or long lived
type AccountToken struct {
	// The token to be used for accessing resources
	AccessToken string `json:"access_token"`
	// RefreshToken to be used to generate a new token
	RefreshToken string `json:"refresh_token"`
	// Time of token creation
	Created time.Time `json:"created"`
	// Time of token expiry
	Expiry time.Time `json:"expiry"`
}

// Expired returns a boolean indicating if the token needs to be refreshed
func (t *AccountToken) Expired() bool {
	return t.Expiry.Unix() < time.Now().Unix()
}

// Resource is an entity such as a user or
type Resource struct {
	// Name of the resource, e.g. go.micro.service.notes
	Name string `json:"name"`
	// Type of resource, e.g. service
	Type string `json:"type"`
	// Endpoint resource e.g NotesService.Create
	Endpoint string `json:"endpoint"`
}

// Access defines the type of access a rule grants
type Access int

const (
	// AccessGranted to a resource
	AccessGranted Access = iota
	// AccessDenied to a resource
	AccessDenied
)

// Rule is used to verify access to a resource
type Rule struct {
	// ID of the rule, e.g. "public"
	ID string
	// Scope the rule requires, a blank scope indicates open to the public and * indicates the rule
	// applies to any valid account
	Scope string
	// Resource the rule applies to
	Resource *Resource
	// Access determines if the rule grants or denies access to the resource
	Access Access
	// Priority the rule should take when verifying a request, the higher the value the sooner the
	// rule will be applied
	Priority int32
}

// Auth provides authentication and authorization
type Auth interface {
	// Init the auth
	Init(opts ...Option)
	// Options set for auth
	Options() Options
	// Generate a new account
	Generate(id string, opts ...GenerateOption) (*Account, error)
	// Verify an account has access to a resource using the rules
	Verify(acc *Account, res *Resource, opts ...VerifyOption) error
	// Inspect a token
	Inspect(token string) (*Account, error)
	// Token generated using refresh token or credentials
	Token(opts ...TokenOption) (*AccountToken, error)
	// Grant access to a resource
	Grant(rule *Rule) error
	// Revoke access to a resource
	Revoke(rule *Rule) error
	// Rules returns all the rules used to verify requests
	Rules(...RulesOption) ([]*Rule, error)
	// String returns the name of the implementation
	String() string
}

// Generate a new account
func Generate(id string, opts ...GenerateOption) (*Account, error) {
	return DefaultAuth.Generate(id, opts...)
}

// Verify an account has access to a resource using the rules
func Verify(acc *Account, res *Resource, opts ...VerifyOption) error {
	return DefaultAuth.Verify(acc, res, opts...)
}

// Inspect a token
func Inspect(token string) (*Account, error) {
	return DefaultAuth.Inspect(token)
}

// Token generated using refresh token or credentials
func Token(opts ...TokenOption) (*AccountToken, error) {
	return DefaultAuth.Token(opts...)
}

// Grant access to a resource
func Grant(rule *Rule) error {
	return DefaultAuth.Grant(rule)
}

// Revoke access to a resource
func Revoke(rule *Rule) error {
	return DefaultAuth.Revoke(rule)
}

// Rules returns all the rules used to verify requests
func Rules(...RulesOption) ([]*Rule, error) {
	return DefaultAuth.Rules()
}

type accountKey struct{}

// AccountFromContext gets the account from the context, which
// is set by the auth wrapper at the start of a call. If the account
// is not set, a nil account will be returned. The error is only returned
// when there was a problem retrieving an account
func AccountFromContext(ctx context.Context) (*Account, bool) {
	acc, ok := ctx.Value(accountKey{}).(*Account)
	return acc, ok
}

// ContextWithAccount sets the account in the context
func ContextWithAccount(ctx context.Context, account *Account) context.Context {
	return context.WithValue(ctx, accountKey{}, account)
}
<|endoftext|>


<commit_before>package service

import (
	"fmt"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/utils/featureflag"
	"github.com/juju/utils/shell"

	"github.com/juju/juju/feature"
	"github.com/juju/juju/service/common"
	"github.com/juju/juju/service/systemd"
	"github.com/juju/juju/service/upstart"
	"github.com/juju/juju/service/windows"
	"github.com/juju/juju/version"
)

// This exists to allow patching during tests.
var getVersion = func() version.Binary {
	return version.Current
}

// DiscoverService returns an interface to a service apropriate
// for the current system
func DiscoverService(name string, conf common.Conf) (Service, error) {
	initName, err := discoverInitSystem()
	if err != nil {
		return nil, errors.Trace(err)
	}

	service, err := NewService(name, conf, initName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return service, nil
}

func discoverInitSystem() (string, error) {
	initName, err := discoverLocalInitSystem()
	if errors.IsNotFound(err) {
		// Fall back to checking the juju version.
		jujuVersion := getVersion()
		versionInitName, ok := VersionInitSystem(jujuVersion)
		if !ok {
			// The key error is the one from discoverLocalInitSystem so
			// that is what we return.
			return "", errors.Trace(err)
		}
		initName = versionInitName
	} else if err != nil {
		return "", errors.Trace(err)
	}
	return initName, nil
}

// VersionInitSystem returns an init system name based on the provided
// version info. If one cannot be identified then false if returned
// for the second return value.
func VersionInitSystem(vers version.Binary) (string, bool) {
	initName, ok := versionInitSystem(vers)
	if !ok {
		logger.Errorf("could not identify init system from juju version info (%#v)", vers)
		return "", false
	}
	logger.Debugf("discovered init system %q from juju version info (%#v)", initName, vers)
	return initName, true
}

func versionInitSystem(vers version.Binary) (string, bool) {
	switch vers.OS {
	case version.Windows:
		return InitSystemWindows, true
	case version.Ubuntu:
		switch vers.Series {
		case "precise", "quantal", "raring", "saucy", "trusty", "utopic":
			return InitSystemUpstart, true
		case "":
			return "", false
		default:
			// Check for pre-precise releases.
			os, _ := version.GetOSFromSeries(vers.Series)
			if os == version.Unknown {
				return "", false
			}
			// vivid and later
			if featureflag.Enabled(feature.LegacyUpstart) {
				return InitSystemUpstart, true
			}
			return InitSystemSystemd, true
		}
		// TODO(ericsnow) Support other OSes, like version.CentOS.
	default:
		return "", false
	}
}

var discoveryFuncs = map[string]func() (bool, error){
	InitSystemSystemd: systemd.IsRunning,
	InitSystemUpstart: upstart.IsRunning,
	InitSystemWindows: windows.IsRunning,
}

func discoverLocalInitSystem() (string, error) {
	for initName, isLocal := range discoveryFuncs {
		local, err := isLocal()
		if err != nil {
			return "", errors.Trace(err)
		}
		if local {
			logger.Debugf("discovered init system %q from local host", initName)
			return initName, nil
		}
	}
	return "", errors.NotFoundf("init system (based on local host)")
}

const discoverInitSystemScript = `
# Use guaranteed discovery mechanisms for known init systems.
if [[ -d /run/systemd/system ]]; then
    echo -n systemd
    exit 0
elif /sbin/initctl --system list 2>&1 > /dev/null; then
    echo -n upstart
    exit 0
fi

# uh-oh
exit 1
`

// DiscoverInitSystemScript returns the shell script to use when
// discovering the local init system. The script is quite specific to
// bash, so it includes an explicit bash shbang.
func DiscoverInitSystemScript() string {
	renderer := shell.BashRenderer{}
	data := renderer.RenderScript([]string{discoverInitSystemScript})
	return string(data)
}

// shellCase is the template for a bash case statement, for use in
// newShellSelectCommand.
const shellCase = `
case "$%s" in
%s
*)
    %s
    ;;
esac`

// newShellSelectCommand creates a bash case statement with clause for
// each of the linux init systems. The body of each clause comes from
// calling the provided handler with the init system name. If the
// handler does not support the args then it returns a false "ok" value.
func newShellSelectCommand(envVarName, dflt string, handler func(string) (string, bool)) string {
	var cases []string
	for _, initSystem := range linuxInitSystems {
		cmd, ok := handler(initSystem)
		if !ok {
			continue
		}
		cases = append(cases, initSystem+")", "    "+cmd, "    ;;")
	}
	if len(cases) == 0 {
		return ""
	}

	return fmt.Sprintf(shellCase[1:], envVarName, strings.Join(cases, "\n"), dflt)
}
<commit_msg>Fix bash script.<commit_after>package service

import (
	"fmt"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/utils/featureflag"
	"github.com/juju/utils/shell"

	"github.com/juju/juju/feature"
	"github.com/juju/juju/service/common"
	"github.com/juju/juju/service/systemd"
	"github.com/juju/juju/service/upstart"
	"github.com/juju/juju/service/windows"
	"github.com/juju/juju/version"
)

// This exists to allow patching during tests.
var getVersion = func() version.Binary {
	return version.Current
}

// DiscoverService returns an interface to a service apropriate
// for the current system
func DiscoverService(name string, conf common.Conf) (Service, error) {
	initName, err := discoverInitSystem()
	if err != nil {
		return nil, errors.Trace(err)
	}

	service, err := NewService(name, conf, initName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return service, nil
}

func discoverInitSystem() (string, error) {
	initName, err := discoverLocalInitSystem()
	if errors.IsNotFound(err) {
		// Fall back to checking the juju version.
		jujuVersion := getVersion()
		versionInitName, ok := VersionInitSystem(jujuVersion)
		if !ok {
			// The key error is the one from discoverLocalInitSystem so
			// that is what we return.
			return "", errors.Trace(err)
		}
		initName = versionInitName
	} else if err != nil {
		return "", errors.Trace(err)
	}
	return initName, nil
}

// VersionInitSystem returns an init system name based on the provided
// version info. If one cannot be identified then false if returned
// for the second return value.
func VersionInitSystem(vers version.Binary) (string, bool) {
	initName, ok := versionInitSystem(vers)
	if !ok {
		logger.Errorf("could not identify init system from juju version info (%#v)", vers)
		return "", false
	}
	logger.Debugf("discovered init system %q from juju version info (%#v)", initName, vers)
	return initName, true
}

func versionInitSystem(vers version.Binary) (string, bool) {
	switch vers.OS {
	case version.Windows:
		return InitSystemWindows, true
	case version.Ubuntu:
		switch vers.Series {
		case "precise", "quantal", "raring", "saucy", "trusty", "utopic":
			return InitSystemUpstart, true
		case "":
			return "", false
		default:
			// Check for pre-precise releases.
			os, _ := version.GetOSFromSeries(vers.Series)
			if os == version.Unknown {
				return "", false
			}
			// vivid and later
			if featureflag.Enabled(feature.LegacyUpstart) {
				return InitSystemUpstart, true
			}
			return InitSystemSystemd, true
		}
		// TODO(ericsnow) Support other OSes, like version.CentOS.
	default:
		return "", false
	}
}

var discoveryFuncs = map[string]func() (bool, error){
	InitSystemSystemd: systemd.IsRunning,
	InitSystemUpstart: upstart.IsRunning,
	InitSystemWindows: windows.IsRunning,
}

func discoverLocalInitSystem() (string, error) {
	for initName, isLocal := range discoveryFuncs {
		local, err := isLocal()
		if err != nil {
			return "", errors.Trace(err)
		}
		if local {
			logger.Debugf("discovered init system %q from local host", initName)
			return initName, nil
		}
	}
	return "", errors.NotFoundf("init system (based on local host)")
}

const discoverInitSystemScript = `
# Use guaranteed discovery mechanisms for known init systems.
if [[ -d /run/systemd/system ]]; then
    echo -n systemd
    exit 0
elif [[ -f /sbin/initctl ]] && /sbin/initctl --system list 2>&1 > /dev/null; then
    echo -n upstart
    exit 0
fi

# uh-oh
exit 1
`

// DiscoverInitSystemScript returns the shell script to use when
// discovering the local init system. The script is quite specific to
// bash, so it includes an explicit bash shbang.
func DiscoverInitSystemScript() string {
	renderer := shell.BashRenderer{}
	data := renderer.RenderScript([]string{discoverInitSystemScript})
	return string(data)
}

// shellCase is the template for a bash case statement, for use in
// newShellSelectCommand.
const shellCase = `
case "$%s" in
%s
*)
    %s
    ;;
esac`

// newShellSelectCommand creates a bash case statement with clause for
// each of the linux init systems. The body of each clause comes from
// calling the provided handler with the init system name. If the
// handler does not support the args then it returns a false "ok" value.
func newShellSelectCommand(envVarName, dflt string, handler func(string) (string, bool)) string {
	var cases []string
	for _, initSystem := range linuxInitSystems {
		cmd, ok := handler(initSystem)
		if !ok {
			continue
		}
		cases = append(cases, initSystem+")", "    "+cmd, "    ;;")
	}
	if len(cases) == 0 {
		return ""
	}

	return fmt.Sprintf(shellCase[1:], envVarName, strings.Join(cases, "\n"), dflt)
}
<|endoftext|>


<commit_before>// Copyright 2017 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 format

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.8.0"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(color.Error, chr)
		}
	}

	fmt.Fprintln(color.Error)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(color.Error, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(color.Error, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(color.Error, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(color.Error, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(color.Error, ", ")
		}
		num++
	}
	fmt.Fprintln(color.Error)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(color.Error)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}

		fmt.Fprintf(color.Error, "%s%s %s %s\n",
			blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)

			fmt.Fprintf(color.Error, "%s%s %s\n",
				yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(color.Error, " ")
		}
	}
	r.Fprintln(color.Error, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(color.Error, Version)
	pad(rightmost - len(Author))
	y.Fprintln(color.Error, Author)
	pad(rightmost - len(Description))
	y.Fprintf(color.Error, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<commit_msg>v3.8.1 release<commit_after>// Copyright 2017 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 format

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.8.1"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(color.Error, chr)
		}
	}

	fmt.Fprintln(color.Error)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(color.Error, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(color.Error, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(color.Error, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(color.Error, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(color.Error, ", ")
		}
		num++
	}
	fmt.Fprintln(color.Error)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(color.Error)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}

		fmt.Fprintf(color.Error, "%s%s %s %s\n",
			blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)

			fmt.Fprintf(color.Error, "%s%s %s\n",
				yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(color.Error, " ")
		}
	}
	r.Fprintln(color.Error, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(color.Error, Version)
	pad(rightmost - len(Author))
	y.Fprintln(color.Error, Author)
	pad(rightmost - len(Description))
	y.Fprintf(color.Error, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<|endoftext|>


<commit_before>// Copyright 2017 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 format

import (
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.8.2"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		if addr.Netblock == nil {
			continue
		}

		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	FprintEnumerationSummary(color.Error, total, tags, asns, demo)
}

// FprintEnumerationSummary outputs the summary information utilized by the command-line tools.
func FprintEnumerationSummary(out io.Writer, total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(out, chr)
		}
	}

	fmt.Fprintln(out)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(out, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(out, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(out, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(out, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(out, ", ")
		}
		num++
	}
	fmt.Fprintln(out)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(out)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}
		fmt.Fprintf(out, "%s%s %s %s\n", blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)
			fmt.Fprintf(out, "%s%s %s\n", yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	FprintBanner(color.Error)
}

// FprintBanner outputs the Amass banner the same for all tools.
func FprintBanner(out io.Writer) {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(out, " ")
		}
	}
	r.Fprintln(out, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(out, Version)
	pad(rightmost - len(Author))
	y.Fprintln(out, Author)
	pad(rightmost - len(Description))
	y.Fprintf(out, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<commit_msg>v3.9.0 release<commit_after>// Copyright 2017 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 format

import (
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.9.0"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		if addr.Netblock == nil {
			continue
		}

		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	FprintEnumerationSummary(color.Error, total, tags, asns, demo)
}

// FprintEnumerationSummary outputs the summary information utilized by the command-line tools.
func FprintEnumerationSummary(out io.Writer, total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(out, chr)
		}
	}

	fmt.Fprintln(out)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(out, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(out, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(out, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(out, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(out, ", ")
		}
		num++
	}
	fmt.Fprintln(out)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(out)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}
		fmt.Fprintf(out, "%s%s %s %s\n", blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)
			fmt.Fprintf(out, "%s%s %s\n", yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	FprintBanner(color.Error)
}

// FprintBanner outputs the Amass banner the same for all tools.
func FprintBanner(out io.Writer) {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(out, " ")
		}
	}
	r.Fprintln(out, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(out, Version)
	pad(rightmost - len(Author))
	y.Fprintln(out, Author)
	pad(rightmost - len(Description))
	y.Fprintf(out, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Andreas Auernhammer. All rights reserved.
// Use of this source code is governed by a license that can be
// found in the LICENSE file.

// Package chacha implements some low-level functions of the
// ChaCha cipher family.
package chacha // import "github.com/aead/chacha20/chacha"

import "encoding/binary"

// NonceSize is the size of the ChaCha20 nonce in bytes.
const NonceSize = 8

// INonceSize is the size of the IETF-ChaCha20 nonce in bytes.
const INonceSize = 12

// XNonceSize is the size of the XChaCha20 nonce in bytes.
const XNonceSize = 24

var (
	useSSE2  bool
	useSSSE3 bool
	useAVX2  bool
)

func setup(state *[64]byte, nonce []byte, key *[32]byte) {
	var Nonce [16]byte
	switch len(nonce) {
	case NonceSize:
		copy(Nonce[8:], nonce)
		initialize(state, key, &Nonce)
	case INonceSize:
		copy(Nonce[4:], nonce)
		initialize(state, key, &Nonce)
	case XNonceSize:
		var tmpKey [32]byte
		var hNonce [16]byte

		copy(hNonce[:], nonce[:16])
		hChaCha20(&tmpKey, &hNonce, key)
		copy(Nonce[8:], nonce[16:])
		initialize(state, &tmpKey, &Nonce)

		// BUG(aead): A "good" compiler will remove this (optimizations)
		//			  But using the provided key instead of tmpKey,
		//			  will change the key (-> probably confuses users)
		for i := range tmpKey {
			tmpKey[i] = 0
		}
	default:
		panic("invalid nonce size") // TODO (add error handling)
	}
}

// XORKeyStream crypts bytes from src to dst using the given key, nonce and counter.
// The rounds argument specifies the number of rounds performed for keystream
// generation - valid values are 8, 12 or 20. The src and dst may be the same slice
// but otherwise should not overlap. If len(dst) < len(src) this function panics.
func XORKeyStream(dst, src, nonce []byte, key *[32]byte, rounds int) {
	if rounds != 20 && rounds != 12 && rounds != 8 {
		panic("chacha20/chacha: rounds must be a 8, 12, or 20")
	}
	if len(dst) < len(src) {
		panic("chacha20/chacha: dst buffer is to small")
	}
	if len(nonce) == INonceSize && uint64(len(src)) > (1<<38) {
		panic("chacha20/chacha: src is too large")
	}

	var block, state [64]byte
	setup(&state, nonce, key)
	xorKeyStream(dst, src, &block, &state, rounds)
}

// Cipher implements ChaCha/X for a given number of rounds X.
type Cipher struct {
	state, block [64]byte
	off          int
	rounds       int // 20 for ChaCha20
}

// NewCipher returns a new *chacha.Cipher implementing the ChaCha/X (X = 8, 12 or 20)
// stream cipher. The nonce must be unique for one key for all time.
func NewCipher(nonce []byte, key *[32]byte, rounds int) *Cipher {
	if rounds != 20 && rounds != 12 && rounds != 8 {
		panic("chacha20/chacha: rounds must be a 8, 12, or 20")
	}

	c := new(Cipher)
	c.rounds = rounds
	setup(&(c.state), nonce, key)

	return c
}

// XORKeyStream crypts bytes from src to dst. Src and dst may be the same slice
// but otherwise should not overlap. If len(dst) < len(src) the function panics.
func (c *Cipher) XORKeyStream(dst, src []byte) {
	if len(dst) < len(src) {
		panic("chacha20/chacha: dst buffer is to small")
	}

	if c.off > 0 {
		n := len(c.block[c.off:])
		if len(src) <= n {
			for i, v := range src {
				dst[i] = v ^ c.block[c.off]
				c.off++
			}
			if c.off == 64 {
				c.off = 0
			}
			return
		}

		for i, v := range c.block[c.off:] {
			dst[i] = src[i] ^ v
		}
		src = src[n:]
		dst = dst[n:]
		c.off = 0
	}

	c.off += xorKeyStream(dst, src, &(c.block), &(c.state), c.rounds)
}

func (c *Cipher) SetCounter(ctr uint64) {
	binary.LittleEndian.PutUint64(c.state[48:], ctr)
}
<commit_msg>v2.0:  set counter correctly<commit_after>// Copyright (c) 2016 Andreas Auernhammer. All rights reserved.
// Use of this source code is governed by a license that can be
// found in the LICENSE file.

// Package chacha implements some low-level functions of the
// ChaCha cipher family.
package chacha // import "github.com/aead/chacha20/chacha"

import "encoding/binary"

// NonceSize is the size of the ChaCha20 nonce in bytes.
const NonceSize = 8

// INonceSize is the size of the IETF-ChaCha20 nonce in bytes.
const INonceSize = 12

// XNonceSize is the size of the XChaCha20 nonce in bytes.
const XNonceSize = 24

var (
	useSSE2  bool
	useSSSE3 bool
	useAVX2  bool
)

func setup(state *[64]byte, nonce []byte, key *[32]byte) {
	var Nonce [16]byte
	switch len(nonce) {
	case NonceSize:
		copy(Nonce[8:], nonce)
		initialize(state, key, &Nonce)
	case INonceSize:
		copy(Nonce[4:], nonce)
		initialize(state, key, &Nonce)
	case XNonceSize:
		var tmpKey [32]byte
		var hNonce [16]byte

		copy(hNonce[:], nonce[:16])
		hChaCha20(&tmpKey, &hNonce, key)
		copy(Nonce[8:], nonce[16:])
		initialize(state, &tmpKey, &Nonce)

		// BUG(aead): A "good" compiler will remove this (optimizations)
		//			  But using the provided key instead of tmpKey,
		//			  will change the key (-> probably confuses users)
		for i := range tmpKey {
			tmpKey[i] = 0
		}
	default:
		panic("invalid nonce size") // TODO (add error handling)
	}
}

// XORKeyStream crypts bytes from src to dst using the given key, nonce and counter.
// The rounds argument specifies the number of rounds performed for keystream
// generation - valid values are 8, 12 or 20. The src and dst may be the same slice
// but otherwise should not overlap. If len(dst) < len(src) this function panics.
func XORKeyStream(dst, src, nonce []byte, key *[32]byte, rounds int) {
	if rounds != 20 && rounds != 12 && rounds != 8 {
		panic("chacha20/chacha: rounds must be a 8, 12, or 20")
	}
	if len(dst) < len(src) {
		panic("chacha20/chacha: dst buffer is to small")
	}
	if len(nonce) == INonceSize && uint64(len(src)) > (1<<38) {
		panic("chacha20/chacha: src is too large")
	}

	var block, state [64]byte
	setup(&state, nonce, key)
	xorKeyStream(dst, src, &block, &state, rounds)
}

// Cipher implements ChaCha/X for a given number of rounds X.
type Cipher struct {
	state, block [64]byte
	off          int
	rounds       int // 20 for ChaCha20
	noncesize    int
}

// NewCipher returns a new *chacha.Cipher implementing the ChaCha/X (X = 8, 12 or 20)
// stream cipher. The nonce must be unique for one key for all time.
func NewCipher(nonce []byte, key *[32]byte, rounds int) *Cipher {
	if rounds != 20 && rounds != 12 && rounds != 8 {
		panic("chacha20/chacha: rounds must be a 8, 12, or 20")
	}

	c := new(Cipher)
	setup(&(c.state), nonce, key)
	c.rounds = rounds

	if len(nonce) == INonceSize {
		c.noncesize = INonceSize
	} else {
		c.noncesize = NonceSize
	}

	return c
}

// XORKeyStream crypts bytes from src to dst. Src and dst may be the same slice
// but otherwise should not overlap. If len(dst) < len(src) the function panics.
func (c *Cipher) XORKeyStream(dst, src []byte) {
	if len(dst) < len(src) {
		panic("chacha20/chacha: dst buffer is to small")
	}

	if c.off > 0 {
		n := len(c.block[c.off:])
		if len(src) <= n {
			for i, v := range src {
				dst[i] = v ^ c.block[c.off]
				c.off++
			}
			if c.off == 64 {
				c.off = 0
			}
			return
		}

		for i, v := range c.block[c.off:] {
			dst[i] = src[i] ^ v
		}
		src = src[n:]
		dst = dst[n:]
		c.off = 0
	}

	c.off += xorKeyStream(dst, src, &(c.block), &(c.state), c.rounds)
}

func (c *Cipher) SetCounter(ctr uint64) {
	if c.noncesize == INonceSize {
		binary.LittleEndian.PutUint32(c.state[48:], uint32(ctr))
	} else {
		binary.LittleEndian.PutUint64(c.state[48:], ctr)
	}
	c.off = 0
}
<|endoftext|>


<commit_before>package chain

import (
	"math/big"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethutil"
)

/*
 * This is the special genesis block.
 */

var ZeroHash256 = make([]byte, 32)
var ZeroHash160 = make([]byte, 20)
var ZeroHash512 = make([]byte, 64)
var EmptyShaList = crypto.Sha3(ethutil.Encode([]interface{}{}))

var GenesisHeader = []interface{}{
	// Previous hash (none)
	ZeroHash256,
	// Empty uncles
	EmptyShaList,
	// Coinbase
	ZeroHash160,
	// Root state
	EmptyShaList,
	// tx sha
	EmptyShaList,
	// receipt list
	EmptyShaList,
	// bloom
	ZeroHash512,
	// Difficulty
	//ethutil.BigPow(2, 22),
	big.NewInt(131072),
	// Number
	ethutil.Big0,
	// Block minimum gas price
	ethutil.Big0,
	// Block upper gas bound
	big.NewInt(1000000),
	// Block gas used
	ethutil.Big0,
	// Time
	ethutil.Big0,
	// Extra
	nil,
	// Nonce
	crypto.Sha3(big.NewInt(42).Bytes()),
}

var Genesis = []interface{}{GenesisHeader, []interface{}{}, []interface{}{}}
<commit_msg>Fixed genesis<commit_after>package chain

import (
	"math/big"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethutil"
)

/*
 * This is the special genesis block.
 */

var ZeroHash256 = make([]byte, 32)
var ZeroHash160 = make([]byte, 20)
var ZeroHash512 = make([]byte, 64)
var EmptyShaList = crypto.Sha3(ethutil.Encode([]interface{}{}))
var EmptyListRoot = crypto.Sha3(ethutil.Encode(""))

var GenesisHeader = []interface{}{
	// Previous hash (none)
	ZeroHash256,
	// Empty uncles
	EmptyShaList,
	// Coinbase
	ZeroHash160,
	// Root state
	EmptyShaList,
	// tx root
	EmptyListRoot,
	// receipt root
	EmptyListRoot,
	// bloom
	ZeroHash512,
	// Difficulty
	//ethutil.BigPow(2, 22),
	big.NewInt(131072),
	// Number
	ethutil.Big0,
	// Block minimum gas price
	ethutil.Big0,
	// Block upper gas bound
	big.NewInt(1000000),
	// Block gas used
	ethutil.Big0,
	// Time
	ethutil.Big0,
	// Extra
	nil,
	// Nonce
	crypto.Sha3(big.NewInt(42).Bytes()),
}

var Genesis = []interface{}{GenesisHeader, []interface{}{}, []interface{}{}}
<|endoftext|>


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

package aerospike

import (
	"fmt"
	"reflect"
	"runtime"
	"sync"

	. "github.com/aerospike/aerospike-client-go/types"
	. "github.com/aerospike/aerospike-client-go/types/atomic"
)

type Result struct {
	Record *Record
	Err    error
}

// String implements the Stringer interface
func (res *Result) String() string {
	if res.Record != nil {
		return fmt.Sprintf("%v", res.Record)
	}
	return fmt.Sprintf("%v", res.Err)
}

// Objectset encapsulates the result of Scan and Query commands.
type objectset struct {
	// a reference to the object channel to close on end signal
	objChan reflect.Value

	// Errors is a channel on which all errors will be sent back.
	// NOTE: Do not use Errors directly. Range on channel returned by Results() instead.
	// This field is deprecated and will be unexported in the future
	Errors chan error

	wgGoroutines sync.WaitGroup
	goroutines   *AtomicInt

	active    *AtomicBool
	cancelled chan struct{}

	chanLock sync.Mutex

	taskId uint64
}

// TaskId returns the transactionId/jobId sent to the server for this recordset.
func (os *objectset) TaskId() uint64 {
	return os.taskId
}

// Recordset encapsulates the result of Scan and Query commands.
type Recordset struct {
	objectset

	// Records is a channel on which the resulting records will be sent back.
	// NOTE: Do not use Records directly. Range on channel returned by Results() instead.
	// Will be unexported in the future
	Records chan *Record
}

// makes sure the recordset is closed eventually, even if it is not consumed
func recordsetFinalizer(rs *Recordset) {
	rs.Close()
}

// newObjectset generates a new RecordSet instance.
func newObjectset(objChan reflect.Value, goroutines int, taskId uint64) *objectset {

	if objChan.Kind() != reflect.Chan ||
		objChan.Type().Elem().Kind() != reflect.Ptr ||
		objChan.Type().Elem().Elem().Kind() != reflect.Struct {
		panic("Scan/Query object channels should be of type `chan *T`")
	}

	rs := &objectset{
		objChan:    objChan,
		Errors:     make(chan error, goroutines),
		active:     NewAtomicBool(true),
		goroutines: NewAtomicInt(goroutines),
		cancelled:  make(chan struct{}),
		taskId:     taskId,
	}
	rs.wgGoroutines.Add(goroutines)

	return rs
}

// newRecordset generates a new RecordSet instance.
func newRecordset(recSize, goroutines int, taskId uint64) *Recordset {
	var nilChan chan *struct{}

	rs := &Recordset{
		Records:   make(chan *Record, recSize),
		objectset: *newObjectset(reflect.ValueOf(nilChan), goroutines, taskId),
	}

	runtime.SetFinalizer(rs, recordsetFinalizer)
	return rs
}

// IsActive returns true if the operation hasn't been finished or cancelled.
func (rcs *Recordset) IsActive() bool {
	return rcs.active.Get()
}

// Read reads the next record from the Recordset. If the Recordset has been
// closed, it returns ErrRecordsetClosed.
func (rcs *Recordset) Read() (record *Record, err error) {
	var ok bool

L:
	select {
	case record, ok = <-rcs.Records:
		if !ok {
			err = ErrRecordsetClosed
		}
	case err = <-rcs.Errors:
		if err == nil {
			// if err == nil, it means the Errors chan has been closed
			// we should not return nil as an error, so we should listen
			// to other chans again to determine either cancellation,
			// or normal EOR
			goto L
		}
	}

	return record, err
}

// Results returns a new receive-only channel with the results of the Scan/Query.
// This is a more idiomatic approach to the iterator pattern in getting the
// results back from the recordset, and doesn't require the user to write the
// ugly select in their code.
// Result contains a Record and an error reference.
//
// Example:
//
//  recordset, err := client.ScanAll(nil, namespace, set)
//  handleError(err)
//  for res := range recordset.Results() {
//    if res.Err != nil {
//      // handle error here
//    } else {
//      // process record here
//      fmt.Println(res.Record.Bins)
//    }
//  }
func (rcs *Recordset) Results() <-chan *Result {
	recCap := cap(rcs.Records)
	if recCap < 1 {
		recCap = 1
	}
	res := make(chan *Result, recCap)

	select {
	case <-rcs.cancelled:
		// Bail early and give the caller a channel for nothing -- it's
		// functionally wasted memory, but the caller did something
		// after close, so it's their own doing.
		close(res)
		return res
	default:
	}

	go func(cancelled <-chan struct{}) {
		defer close(res)
		for {
			record, err := rcs.Read()
			if err == ErrRecordsetClosed {
				return
			}

			result := &Result{Record: record, Err: err}
			select {
			case <-cancelled:
				return
			case res <- result:

			}
		}
	}(rcs.cancelled)

	return res
}

// Close all streams from different nodes. A successful close return nil,
// subsequent calls to the method will return ErrRecordsetClosed.
func (rcs *Recordset) Close() error {
	// check and close cancelled channel
	// this will broadcast to all commands listening to the channel
	select {
	case <-rcs.cancelled:
	default:
		close(rcs.cancelled)
	}

	// do it only once
	if !rcs.active.CompareAndToggle(true) {
		return ErrRecordsetClosed
	}

	// wait till all goroutines are done, and signalEnd is called by the scan command
	rcs.wgGoroutines.Wait()

	return nil
}

func (rcs *Recordset) signalEnd() {
	rcs.wgGoroutines.Done()
	if rcs.goroutines.DecrementAndGet() == 0 {
		// mark the recordset as closed
		rcs.active.Set(false)

		rcs.chanLock.Lock()
		defer rcs.chanLock.Unlock()

		if rcs.Records != nil {
			close(rcs.Records)
		} else if rcs.objChan.IsValid() {
			rcs.objChan.Close()
		}

		close(rcs.Errors)
	}
}

func (rcs *Recordset) sendError(err error) {
	rcs.chanLock.Lock()
	defer rcs.chanLock.Unlock()
	if rcs.IsActive() {
		rcs.Errors <- err
	}
}
<commit_msg>Use a secondary atomic to close the cancelled channel exactly once<commit_after>// Copyright 2013-2016 Aerospike, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aerospike

import (
	"fmt"
	"reflect"
	"runtime"
	"sync"

	. "github.com/aerospike/aerospike-client-go/types"
	. "github.com/aerospike/aerospike-client-go/types/atomic"
)

type Result struct {
	Record *Record
	Err    error
}

// String implements the Stringer interface
func (res *Result) String() string {
	if res.Record != nil {
		return fmt.Sprintf("%v", res.Record)
	}
	return fmt.Sprintf("%v", res.Err)
}

// Objectset encapsulates the result of Scan and Query commands.
type objectset struct {
	// a reference to the object channel to close on end signal
	objChan reflect.Value

	// Errors is a channel on which all errors will be sent back.
	// NOTE: Do not use Errors directly. Range on channel returned by Results() instead.
	// This field is deprecated and will be unexported in the future
	Errors chan error

	wgGoroutines sync.WaitGroup
	goroutines   *AtomicInt

	closed, active *AtomicBool
	cancelled      chan struct{}

	chanLock sync.Mutex

	taskId uint64
}

// TaskId returns the transactionId/jobId sent to the server for this recordset.
func (os *objectset) TaskId() uint64 {
	return os.taskId
}

// Recordset encapsulates the result of Scan and Query commands.
type Recordset struct {
	objectset

	// Records is a channel on which the resulting records will be sent back.
	// NOTE: Do not use Records directly. Range on channel returned by Results() instead.
	// Will be unexported in the future
	Records chan *Record
}

// makes sure the recordset is closed eventually, even if it is not consumed
func recordsetFinalizer(rs *Recordset) {
	rs.Close()
}

// newObjectset generates a new RecordSet instance.
func newObjectset(objChan reflect.Value, goroutines int, taskId uint64) *objectset {

	if objChan.Kind() != reflect.Chan ||
		objChan.Type().Elem().Kind() != reflect.Ptr ||
		objChan.Type().Elem().Elem().Kind() != reflect.Struct {
		panic("Scan/Query object channels should be of type `chan *T`")
	}

	rs := &objectset{
		objChan:    objChan,
		Errors:     make(chan error, goroutines),
		active:     NewAtomicBool(true),
		closed:     NewAtomicBool(false),
		goroutines: NewAtomicInt(goroutines),
		cancelled:  make(chan struct{}),
		taskId:     taskId,
	}
	rs.wgGoroutines.Add(goroutines)

	return rs
}

// newRecordset generates a new RecordSet instance.
func newRecordset(recSize, goroutines int, taskId uint64) *Recordset {
	var nilChan chan *struct{}

	rs := &Recordset{
		Records:   make(chan *Record, recSize),
		objectset: *newObjectset(reflect.ValueOf(nilChan), goroutines, taskId),
	}

	runtime.SetFinalizer(rs, recordsetFinalizer)
	return rs
}

// IsActive returns true if the operation hasn't been finished or cancelled.
func (rcs *Recordset) IsActive() bool {
	return rcs.active.Get()
}

// Read reads the next record from the Recordset. If the Recordset has been
// closed, it returns ErrRecordsetClosed.
func (rcs *Recordset) Read() (record *Record, err error) {
	var ok bool

L:
	select {
	case record, ok = <-rcs.Records:
		if !ok {
			err = ErrRecordsetClosed
		}
	case err = <-rcs.Errors:
		if err == nil {
			// if err == nil, it means the Errors chan has been closed
			// we should not return nil as an error, so we should listen
			// to other chans again to determine either cancellation,
			// or normal EOR
			goto L
		}
	}

	return record, err
}

// Results returns a new receive-only channel with the results of the Scan/Query.
// This is a more idiomatic approach to the iterator pattern in getting the
// results back from the recordset, and doesn't require the user to write the
// ugly select in their code.
// Result contains a Record and an error reference.
//
// Example:
//
//  recordset, err := client.ScanAll(nil, namespace, set)
//  handleError(err)
//  for res := range recordset.Results() {
//    if res.Err != nil {
//      // handle error here
//    } else {
//      // process record here
//      fmt.Println(res.Record.Bins)
//    }
//  }
func (rcs *Recordset) Results() <-chan *Result {
	recCap := cap(rcs.Records)
	if recCap < 1 {
		recCap = 1
	}
	res := make(chan *Result, recCap)

	select {
	case <-rcs.cancelled:
		// Bail early and give the caller a channel for nothing -- it's
		// functionally wasted memory, but the caller did something
		// after close, so it's their own doing.
		close(res)
		return res
	default:
	}

	go func(cancelled <-chan struct{}) {
		defer close(res)
		for {
			record, err := rcs.Read()
			if err == ErrRecordsetClosed {
				return
			}

			result := &Result{Record: record, Err: err}
			select {
			case <-cancelled:
				return
			case res <- result:

			}
		}
	}(rcs.cancelled)

	return res
}

// Close all streams from different nodes. A successful close return nil,
// subsequent calls to the method will return ErrRecordsetClosed.
func (rcs *Recordset) Close() error {
	// do it only once
	if !rcs.closed.CompareAndToggle(false) {
		return ErrRecordsetClosed
	}

	// mark the recordset as inactive
	rcs.active.Set(false)

	close(rcs.cancelled)

	// wait till all goroutines are done, and signalEnd is called by the scan command
	rcs.wgGoroutines.Wait()

	return nil
}

func (rcs *Recordset) signalEnd() {
	rcs.wgGoroutines.Done()
	if rcs.goroutines.DecrementAndGet() == 0 {
		// mark the recordset as inactive
		rcs.active.Set(false)

		rcs.chanLock.Lock()
		defer rcs.chanLock.Unlock()

		if rcs.Records != nil {
			close(rcs.Records)
		} else if rcs.objChan.IsValid() {
			rcs.objChan.Close()
		}

		close(rcs.Errors)
	}
}

func (rcs *Recordset) sendError(err error) {
	rcs.chanLock.Lock()
	defer rcs.chanLock.Unlock()
	if rcs.IsActive() {
		rcs.Errors <- err
	}
}
<|endoftext|>


<commit_before>package headlessChrome

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/integrii/interactive"
)

// Debug enables debug output for this package to console
var Debug bool

// ChromePath is the command to execute chrome
var ChromePath = `/Applications/Google Chrome.app/Contents/MacOS/Google Chrome`

// Args are the args that will be used to start chrome
var Args = []string{
	"--headless",
	"--disable-gpu",
	"--repl",
	// "--dump-dom",
	// "--window-size=1024,768",
	// "--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
	// "--verbose",
}

const expectedFirstLine = `Type a Javascript expression to evaluate or "quit" to exit.`
const promptPrefix = `>>>`

// OutputSanitizer puts output coming from the consolw that
// does not begin with the input prompt into the session
// output channel
func (cs *ChromeSession) OutputSanitizer() {
	for text := range cs.Session.Output {
		if !strings.HasPrefix(text, promptPrefix) {
			cs.Output <- text
		}
	}
}

// ChromeSession is an interactive console Session with a Chrome
// instance.
type ChromeSession struct {
	Session *interactive.Session
	Output  chan string
	Input   chan string
}

// Exit exits the running command out by ossuing a 'quit'
// to the chrome console
func (cs *ChromeSession) Exit() {
	cs.Session.Write(`quit`)
	cs.Session.Exit()
}

// Write writes to the Session
func (cs *ChromeSession) Write(s string) {
	if Debug {
		fmt.Println("Writing to console:")
		fmt.Println(s)
	}
	cs.Session.Write(s)
}

// OutputPrinter prints all outputs from the output channel to the cli
func (cs *ChromeSession) OutputPrinter() {
	for l := range cs.Session.Output {
		fmt.Println(l)
	}
}

// forceClose issues a force kill to the command
func (cs *ChromeSession) forceClose() {
	cs.Session.ForceClose()
}

// ClickSelector calls a click() on the supplied selector
func (cs *ChromeSession) ClickSelector(s string) {
	cs.Write(`document.querySelector("` + s + `").click()`)
}

// ClickItemWithInnerHTML clicks an item that has the matching inner html
func (cs *ChromeSession) ClickItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `].click()`)
}

// GetItemWithInnerHTML fetches the item with the specified innerHTML content
func (cs *ChromeSession) GetItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `]`)
}

// GetContentOfItemWithClasses fetches the content of the element with the specified classes
func (cs *ChromeSession) GetContentOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML`)
}

// GetValueOfItemWithClasses returns the form value of the specified item
func (cs *ChromeSession) GetValueOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value`)
}

// GetContentOfItemWithSelector gets the content of an element with the specified selector
func (cs *ChromeSession) GetContentOfItemWithSelector(selector string) {
	cs.Write(`document.querySelector("` + selector + `").innerHTML()`)
}

// ClickItemWithClasses clicks on the first item it finds with the provided classes.
// Multiple classes are separated by spaces
func (cs *ChromeSession) ClickItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].click()`)
}

// SetTextByID sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByID(id string, text string) {
	cs.Write(`document.getElementById("` + id + `").innerHTML = "` + text + `"`)
}

// ClickItemWithID clicks an item with the specified id
func (cs *ChromeSession) ClickItemWithID(id string) {
	cs.Write(`document.getElementById("` + id + `").click()`)
}

// SetTextByClasses sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML = "` + text + `"`)
}

// SetInputTextByClasses sets the input text for an input field
func (cs *ChromeSession) SetInputTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value = "` + text + `"`)
}

// NewBrowser starts a new chrome headless Session.
func NewBrowser(url string) (*ChromeSession, error) {
	var err error

	chromeSession := ChromeSession{}
	chromeSession.Output = make(chan string, 5000)

	// add url as last arg and create new Session
	args := append(Args, url)
	chromeSession.Session, err = interactive.NewSession(ChromePath, args)

	// map output and input channels for easy use
	chromeSession.Input = chromeSession.Session.Input

	go chromeSession.OutputSanitizer()

	firstLine := <-chromeSession.Output
	if !strings.Contains(firstLine, expectedFirstLine) {
		log.Println("ERROR: Unespected first line when initializing headless Chrome console:", firstLine)
	}

	return &chromeSession, err
}
<commit_msg>NewBrowserWithTimeout func that kills the browser after the specified time<commit_after>package headlessChrome

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/integrii/interactive"
)

// Debug enables debug output for this package to console
var Debug bool

// ChromePath is the command to execute chrome
var ChromePath = `/Applications/Google Chrome.app/Contents/MacOS/Google Chrome`

// Args are the args that will be used to start chrome
var Args = []string{
	"--headless",
	"--disable-gpu",
	"--repl",
	// "--dump-dom",
	// "--window-size=1024,768",
	// "--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
	// "--verbose",
}

const expectedFirstLine = `Type a Javascript expression to evaluate or "quit" to exit.`
const promptPrefix = `>>>`

// OutputSanitizer puts output coming from the consolw that
// does not begin with the input prompt into the session
// output channel
func (cs *ChromeSession) OutputSanitizer() {
	for text := range cs.Session.Output {
		if !strings.HasPrefix(text, promptPrefix) {
			cs.Output <- text
		}
	}
}

// ChromeSession is an interactive console Session with a Chrome
// instance.
type ChromeSession struct {
	Session *interactive.Session
	Output  chan string
	Input   chan string
}

// Exit exits the running command out by ossuing a 'quit'
// to the chrome console
func (cs *ChromeSession) Exit() {
	cs.Session.Write(`quit`)
	cs.Session.Exit()
}

// Write writes to the Session
func (cs *ChromeSession) Write(s string) {
	if Debug {
		fmt.Println("Writing to console:")
		fmt.Println(s)
	}
	cs.Session.Write(s)
}

// OutputPrinter prints all outputs from the output channel to the cli
func (cs *ChromeSession) OutputPrinter() {
	for l := range cs.Session.Output {
		fmt.Println(l)
	}
}

// forceClose issues a force kill to the command
func (cs *ChromeSession) forceClose() {
	cs.Session.ForceClose()
}

// ClickSelector calls a click() on the supplied selector
func (cs *ChromeSession) ClickSelector(s string) {
	cs.Write(`document.querySelector("` + s + `").click()`)
}

// ClickItemWithInnerHTML clicks an item that has the matching inner html
func (cs *ChromeSession) ClickItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `].click()`)
}

// GetItemWithInnerHTML fetches the item with the specified innerHTML content
func (cs *ChromeSession) GetItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `]`)
}

// GetContentOfItemWithClasses fetches the content of the element with the specified classes
func (cs *ChromeSession) GetContentOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML`)
}

// GetValueOfItemWithClasses returns the form value of the specified item
func (cs *ChromeSession) GetValueOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value`)
}

// GetContentOfItemWithSelector gets the content of an element with the specified selector
func (cs *ChromeSession) GetContentOfItemWithSelector(selector string) {
	cs.Write(`document.querySelector("` + selector + `").innerHTML()`)
}

// ClickItemWithClasses clicks on the first item it finds with the provided classes.
// Multiple classes are separated by spaces
func (cs *ChromeSession) ClickItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].click()`)
}

// SetTextByID sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByID(id string, text string) {
	cs.Write(`document.getElementById("` + id + `").innerHTML = "` + text + `"`)
}

// ClickItemWithID clicks an item with the specified id
func (cs *ChromeSession) ClickItemWithID(id string) {
	cs.Write(`document.getElementById("` + id + `").click()`)
}

// SetTextByClasses sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML = "` + text + `"`)
}

// SetInputTextByClasses sets the input text for an input field
func (cs *ChromeSession) SetInputTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value = "` + text + `"`)
}

// NewBrowserWithTimeout starts a new chrome headless session
// but limits how long it can run before its killed forcefully.
// A time limit of 0 means there is not a time limit
func NewBrowserWithTimeout(url string, timeout time.Duration) (*ChromeSession, error) {
	var err error

	chromeSession := ChromeSession{}
	chromeSession.Output = make(chan string, 5000)

	// add url as last arg and create new Session
	args := append(Args, url)
	chromeSession.Session, err = interactive.NewSessionWithTimeout(ChromePath, args, timeout)

	// map output and input channels for easy use
	chromeSession.Input = chromeSession.Session.Input

	go chromeSession.OutputSanitizer()

	firstLine := <-chromeSession.Output
	if !strings.Contains(firstLine, expectedFirstLine) {
		log.Println("ERROR: Unespected first line when initializing headless Chrome console:", firstLine)
	}

	return &chromeSession, err
}

// NewBrowser starts a new chrome headless Session.
func NewBrowser(url string) (*ChromeSession, error) {
	return NewBrowserWithTimeout(url, 0)
}
<|endoftext|>


<commit_before>package detour

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/getlantern/golog"
)

type wlEntry struct {
	permanent bool
	addTime   time.Time
}

var (
	log = golog.LoggerFor("detour")
	// if dial or read exceeded this timeout, we consider switch to detour
	timeoutToDetour = 1 * time.Second

	muWhitelist sync.RWMutex
	whitelist   = make(map[string]wlEntry)
)

type dialFunc func(network, addr string) (net.Conn, error)

type detourConn struct {
	muConn sync.RWMutex
	// the actual connection, will change so protect it
	// can't user atomic.Value as the concrete type may vary
	conn net.Conn

	// don't access directly, use inState() and setState() instead
	state uint32

	// the function to dial detour if the site to connect seems blocked
	dialDetour dialFunc

	muBuf sync.Mutex
	// keep track of bytes sent through normal connection
	// so we can resend them when detour
	buf bytes.Buffer

	network, addr string
	readDeadline  time.Time
	writeDeadline time.Time
}

const (
	stateInitial = iota
	stateDirect
	stateDetour
	stateWhitelistCandidate
	stateWhitelist
)

var statesDesc = []string{
	"INITIALLY",
	"DIRECTLY",
	"DETOURED",
	"WHITELIST CANDIDATE",
	"WHITELISTED",
}

// SetTimeout sets the timeout so if dial or read exceeds this timeout, we consider switch to detour
// The value depends on OS and browser and defaults to 1s
// For Windows XP, find TcpMaxConnectRetransmissions in http://support2.microsoft.com/default.aspx?scid=kb;en-us;314053
func SetTimeout(t time.Duration) {
	timeoutToDetour = t
}

func InitWhiteList(wl map[string]time.Time) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	for k, v := range wl {
		whitelist[k] = wlEntry{true, v}
	}
	return
}

func DumpWhiteList() (wl map[string]time.Time) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	for k, v := range whitelist {
		wl[k] = v.addTime
	}
	return
}

func Dialer(dialer dialFunc) dialFunc {
	return func(network, addr string) (conn net.Conn, err error) {
		dc := &detourConn{dialDetour: dialer, network: network, addr: addr}
		if !whitelisted(addr) {
			dc.setState(stateInitial)
			dc.conn, err = net.DialTimeout(network, addr, timeoutToDetour)
			if err == nil {
				log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
				return dc, nil
			}
			log.Debugf("Dial %s to %s failed, try detour: %s", dc.stateDesc(), addr, err)
		}
		dc.setState(stateDetour)
		dc.conn, err = dc.dialDetour(network, addr)
		if err != nil {
			log.Errorf("Dial %s to %s failed", dc.stateDesc(), addr)
			return nil, err
		}
		log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
		return dc, err
	}
}

// Read() implements the function from net.Conn
func (dc *detourConn) Read(b []byte) (n int, err error) {
	conn := dc.getConn()
	if !dc.inState(stateInitial) {
		if n, err = conn.Read(b); err != nil && err != io.EOF {
			log.Tracef("Read from %s %s failed: %s", dc.addr, dc.stateDesc(), err)
			if dc.inState(stateDirect) && blocked(err) {
				log.Tracef("Seems %s still blocked, add to whitelist so will try detour next time", dc.addr)
				addToWl(dc.addr, false)
			} else if wlTemporarily(dc.addr) {
				log.Tracef("Detoured route is still not reliable for %s, not whitelist it", dc.addr)
				removeFromWl(dc.addr)
			}
			return
		}
		log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
		return n, err
	}
	// state will always be settled after first read, safe to clear buffer at end of it
	defer dc.resetBuffer()
	start := time.Now()
	dl := start.Add(timeoutToDetour)
	if !dc.readDeadline.IsZero() && dc.readDeadline.Sub(start) < 2*timeoutToDetour {
		log.Tracef("no time left to test %s, read %s", dc.addr, stateDirect)
		dc.setState(stateDirect)
		return conn.Read(b)
	}
	conn.SetReadDeadline(dl)

	n, err = conn.Read(b)
	conn.SetReadDeadline(dc.readDeadline)
	if err != nil && err != io.EOF {
		ne := fmt.Errorf("Error while read from %s %s, takes %s: %s", dc.addr, dc.stateDesc(), time.Now().Sub(start), err)
		log.Debug(ne)
		if !blocked(err) {
			return n, ne
		}
		return dc.detour(b)
	}
	log.Tracef("Read %d bytes from %s %s, set state to DIRECT", n, dc.addr, dc.stateDesc())
	dc.setState(stateDirect)
	return n, err
}

// Write() implements the function from net.Conn
func (dc *detourConn) Write(b []byte) (n int, err error) {
	if dc.inState(stateInitial) {
		if n, err = dc.writeToBuffer(b); err != nil {
			return n, fmt.Errorf("Unable to write to local buffer: %s", err)
		}
	}
	if n, err = dc.getConn().Write(b); err != nil {
		log.Debugf("Error while write %d bytes to %s %s: %s", len(b), dc.addr, dc.stateDesc(), err)
		return
	}
	log.Tracef("Wrote %d bytes to %s %s", len(b), dc.addr, dc.stateDesc())
	return
}

// Close() implements the function from net.Conn
func (dc *detourConn) Close() error {
	log.Tracef("Closing %s connection to %s", dc.stateDesc(), dc.addr)
	if dc.inState(stateInitial) && wlTemporarily(dc.addr) {
		log.Tracef("no error found till closing, add %s to permanent whitelist", dc.addr)
		addToWl(dc.addr, true)
	}
	return dc.getConn().Close()
}

func (dc *detourConn) LocalAddr() net.Addr {
	return dc.getConn().LocalAddr()
}

func (dc *detourConn) RemoteAddr() net.Addr {
	return dc.getConn().RemoteAddr()
}

func (dc *detourConn) SetDeadline(t time.Time) error {
	dc.SetReadDeadline(t)
	dc.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) SetReadDeadline(t time.Time) error {
	dc.readDeadline = t
	dc.conn.SetReadDeadline(t)
	return nil
}

func (dc *detourConn) SetWriteDeadline(t time.Time) error {
	dc.writeDeadline = t
	dc.conn.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) writeToBuffer(b []byte) (n int, err error) {
	dc.muBuf.Lock()
	n, err = dc.buf.Write(b)
	dc.muBuf.Unlock()
	return
}

func (dc *detourConn) resetBuffer() {
	dc.muBuf.Lock()
	dc.buf.Reset()
	dc.muBuf.Unlock()
}

func (dc *detourConn) detour(b []byte) (n int, err error) {
	if err = dc.setupDetour(); err != nil {
		log.Errorf("Error to setup detour: %s", err)
		return
	}
	if _, err = dc.resend(); err != nil {
		err = fmt.Errorf("Error resend buffer to %s: %s", dc.addr, err)
		log.Error(err)
		return
	}
	// should getConn() again as it has changed
	if n, err = dc.getConn().Read(b); err != nil {
		log.Debugf("Read from %s %s still failed: %s", dc.addr, dc.stateDesc(), err)
		return
	}
	dc.setState(stateDetour)
	addToWl(dc.addr, false)
	log.Tracef("Read %d bytes from %s through detour, set state to %s", n, dc.addr, dc.stateDesc())
	return
}

func (dc *detourConn) resend() (int, error) {
	dc.muBuf.Lock()
	b := dc.buf.Bytes()
	dc.muBuf.Unlock()
	if len(b) > 0 {
		n, err := dc.getConn().Write(b)
		log.Tracef("Resend %d buffered bytes to %s, %d sent", len(b), dc.addr, n)
		return n, err
	}
	return 0, nil
}

func (dc *detourConn) setupDetour() error {
	c, err := dc.dialDetour("tcp", dc.addr)
	if err != nil {
		return err
	}
	log.Tracef("Dialed a new detour connection to %s", dc.addr)
	dc.setConn(c)
	return nil
}

func (dc *detourConn) getConn() (c net.Conn) {
	dc.muConn.RLock()
	defer dc.muConn.RUnlock()
	return dc.conn
}

func (dc *detourConn) setConn(c net.Conn) {
	dc.muConn.Lock()
	oldConn := dc.conn
	dc.conn = c
	dc.muConn.Unlock()
	dc.conn.SetReadDeadline(dc.readDeadline)
	dc.conn.SetWriteDeadline(dc.writeDeadline)
	log.Tracef("Replaced connection to %s from direct to detour and closing old one", dc.addr)
	oldConn.Close()
}

func (dc *detourConn) stateDesc() string {
	return statesDesc[atomic.LoadUint32(&dc.state)]
}

func (dc *detourConn) inState(s uint32) bool {
	return atomic.LoadUint32(&dc.state) == s
}

func (dc *detourConn) setState(s uint32) {
	atomic.StoreUint32(&dc.state, s)
}

func blocked(err error) bool {
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return true
	}
	if oe, ok := err.(*net.OpError); ok && (oe.Err == syscall.EPIPE || oe.Err == syscall.ECONNRESET) {
		return true
	}
	return false
}

func whitelisted(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	_, in := whitelist[addr]
	return in
}

func wlTemporarily(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	p, ok := whitelist[addr]
	return ok && p.permanent == false
}

func addToWl(addr string, permanent bool) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	whitelist[addr] = wlEntry{permanent, time.Now()}
}

func removeFromWl(addr string) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	delete(whitelist, addr)
}
<commit_msg>bug fix<commit_after>package detour

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/getlantern/golog"
)

type wlEntry struct {
	permanent bool
	addTime   time.Time
}

var (
	log = golog.LoggerFor("detour")
	// if dial or read exceeded this timeout, we consider switch to detour
	timeoutToDetour = 1 * time.Second

	muWhitelist sync.RWMutex
	whitelist   = make(map[string]wlEntry)
)

type dialFunc func(network, addr string) (net.Conn, error)

type detourConn struct {
	muConn sync.RWMutex
	// the actual connection, will change so protect it
	// can't user atomic.Value as the concrete type may vary
	conn net.Conn

	// don't access directly, use inState() and setState() instead
	state uint32

	// the function to dial detour if the site to connect seems blocked
	dialDetour dialFunc

	muBuf sync.Mutex
	// keep track of bytes sent through normal connection
	// so we can resend them when detour
	buf bytes.Buffer

	network, addr string
	readDeadline  time.Time
	writeDeadline time.Time
}

const (
	stateInitial = iota
	stateDirect
	stateDetour
	stateWhitelistCandidate
	stateWhitelist
)

var statesDesc = []string{
	"INITIALLY",
	"DIRECTLY",
	"DETOURED",
	"WHITELIST CANDIDATE",
	"WHITELISTED",
}

// SetTimeout sets the timeout so if dial or read exceeds this timeout, we consider switch to detour
// The value depends on OS and browser and defaults to 1s
// For Windows XP, find TcpMaxConnectRetransmissions in http://support2.microsoft.com/default.aspx?scid=kb;en-us;314053
func SetTimeout(t time.Duration) {
	timeoutToDetour = t
}

func InitWhiteList(wl map[string]time.Time) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	for k, v := range wl {
		whitelist[k] = wlEntry{true, v}
	}
	return
}

func DumpWhiteList() (wl map[string]time.Time) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	for k, v := range whitelist {
		wl[k] = v.addTime
	}
	return
}

func Dialer(dialer dialFunc) dialFunc {
	return func(network, addr string) (conn net.Conn, err error) {
		dc := &detourConn{dialDetour: dialer, network: network, addr: addr}
		if !whitelisted(addr) {
			dc.setState(stateInitial)
			dc.conn, err = net.DialTimeout(network, addr, timeoutToDetour)
			if err == nil {
				log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
				return dc, nil
			}
			log.Debugf("Dial %s to %s failed, try detour: %s", dc.stateDesc(), addr, err)
		}
		dc.setState(stateDetour)
		dc.conn, err = dc.dialDetour(network, addr)
		if err != nil {
			log.Errorf("Dial %s to %s failed", dc.stateDesc(), addr)
			return nil, err
		}
		log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
		return dc, err
	}
}

// Read() implements the function from net.Conn
func (dc *detourConn) Read(b []byte) (n int, err error) {
	conn := dc.getConn()
	if !dc.inState(stateInitial) {
		if n, err = conn.Read(b); err != nil && err != io.EOF {
			log.Tracef("Read from %s %s failed: %s", dc.addr, dc.stateDesc(), err)
			if dc.inState(stateDirect) && blocked(err) {
				log.Tracef("Seems %s still blocked, add to whitelist so will try detour next time", dc.addr)
				addToWl(dc.addr, false)
			} else if wlTemporarily(dc.addr) {
				log.Tracef("Detoured route is still not reliable for %s, not whitelist it", dc.addr)
				removeFromWl(dc.addr)
			}
			return
		}
		log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
		return n, err
	}
	// state will always be settled after first read, safe to clear buffer at end of it
	defer dc.resetBuffer()
	start := time.Now()
	dl := start.Add(timeoutToDetour)
	if !dc.readDeadline.IsZero() && dc.readDeadline.Sub(start) < 2*timeoutToDetour {
		log.Tracef("no time left to test %s, read %s", dc.addr, stateDirect)
		dc.setState(stateDirect)
		return conn.Read(b)
	}
	conn.SetReadDeadline(dl)

	n, err = conn.Read(b)
	conn.SetReadDeadline(dc.readDeadline)
	if err != nil && err != io.EOF {
		ne := fmt.Errorf("Error while read from %s %s, takes %s: %s", dc.addr, dc.stateDesc(), time.Now().Sub(start), err)
		log.Debug(ne)
		if !blocked(err) {
			return n, ne
		}
		return dc.detour(b)
	}
	log.Tracef("Read %d bytes from %s %s, set state to DIRECT", n, dc.addr, dc.stateDesc())
	dc.setState(stateDirect)
	return n, err
}

// Write() implements the function from net.Conn
func (dc *detourConn) Write(b []byte) (n int, err error) {
	if dc.inState(stateInitial) {
		if n, err = dc.writeToBuffer(b); err != nil {
			return n, fmt.Errorf("Unable to write to local buffer: %s", err)
		}
	}
	if n, err = dc.getConn().Write(b); err != nil {
		log.Debugf("Error while write %d bytes to %s %s: %s", len(b), dc.addr, dc.stateDesc(), err)
		return
	}
	log.Tracef("Wrote %d bytes to %s %s", len(b), dc.addr, dc.stateDesc())
	return
}

// Close() implements the function from net.Conn
func (dc *detourConn) Close() error {
	log.Tracef("Closing %s connection to %s", dc.stateDesc(), dc.addr)
	if dc.inState(stateDetour) && wlTemporarily(dc.addr) {
		log.Tracef("no error found till closing, add %s to permanent whitelist", dc.addr)
		addToWl(dc.addr, true)
	}
	return dc.getConn().Close()
}

func (dc *detourConn) LocalAddr() net.Addr {
	return dc.getConn().LocalAddr()
}

func (dc *detourConn) RemoteAddr() net.Addr {
	return dc.getConn().RemoteAddr()
}

func (dc *detourConn) SetDeadline(t time.Time) error {
	dc.SetReadDeadline(t)
	dc.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) SetReadDeadline(t time.Time) error {
	dc.readDeadline = t
	dc.conn.SetReadDeadline(t)
	return nil
}

func (dc *detourConn) SetWriteDeadline(t time.Time) error {
	dc.writeDeadline = t
	dc.conn.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) writeToBuffer(b []byte) (n int, err error) {
	dc.muBuf.Lock()
	n, err = dc.buf.Write(b)
	dc.muBuf.Unlock()
	return
}

func (dc *detourConn) resetBuffer() {
	dc.muBuf.Lock()
	dc.buf.Reset()
	dc.muBuf.Unlock()
}

func (dc *detourConn) detour(b []byte) (n int, err error) {
	if err = dc.setupDetour(); err != nil {
		log.Errorf("Error to setup detour: %s", err)
		return
	}
	if _, err = dc.resend(); err != nil {
		err = fmt.Errorf("Error resend buffer to %s: %s", dc.addr, err)
		log.Error(err)
		return
	}
	// should getConn() again as it has changed
	if n, err = dc.getConn().Read(b); err != nil {
		log.Debugf("Read from %s %s still failed: %s", dc.addr, dc.stateDesc(), err)
		return
	}
	dc.setState(stateDetour)
	addToWl(dc.addr, false)
	log.Tracef("Read %d bytes from %s through detour, set state to %s", n, dc.addr, dc.stateDesc())
	return
}

func (dc *detourConn) resend() (int, error) {
	dc.muBuf.Lock()
	b := dc.buf.Bytes()
	dc.muBuf.Unlock()
	if len(b) > 0 {
		n, err := dc.getConn().Write(b)
		log.Tracef("Resend %d buffered bytes to %s, %d sent", len(b), dc.addr, n)
		return n, err
	}
	return 0, nil
}

func (dc *detourConn) setupDetour() error {
	c, err := dc.dialDetour("tcp", dc.addr)
	if err != nil {
		return err
	}
	log.Tracef("Dialed a new detour connection to %s", dc.addr)
	dc.setConn(c)
	return nil
}

func (dc *detourConn) getConn() (c net.Conn) {
	dc.muConn.RLock()
	defer dc.muConn.RUnlock()
	return dc.conn
}

func (dc *detourConn) setConn(c net.Conn) {
	dc.muConn.Lock()
	oldConn := dc.conn
	dc.conn = c
	dc.muConn.Unlock()
	dc.conn.SetReadDeadline(dc.readDeadline)
	dc.conn.SetWriteDeadline(dc.writeDeadline)
	log.Tracef("Replaced connection to %s from direct to detour and closing old one", dc.addr)
	oldConn.Close()
}

func (dc *detourConn) stateDesc() string {
	return statesDesc[atomic.LoadUint32(&dc.state)]
}

func (dc *detourConn) inState(s uint32) bool {
	return atomic.LoadUint32(&dc.state) == s
}

func (dc *detourConn) setState(s uint32) {
	atomic.StoreUint32(&dc.state, s)
}

func blocked(err error) bool {
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return true
	}
	if oe, ok := err.(*net.OpError); ok && (oe.Err == syscall.EPIPE || oe.Err == syscall.ECONNRESET) {
		return true
	}
	return false
}

func whitelisted(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	_, in := whitelist[addr]
	return in
}

func wlTemporarily(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	p, ok := whitelist[addr]
	return ok && p.permanent == false
}

func addToWl(addr string, permanent bool) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	whitelist[addr] = wlEntry{permanent, time.Now()}
}

func removeFromWl(addr string) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	delete(whitelist, addr)
}
<|endoftext|>