<commit_before>package sohop

import (
	"crypto/rand"
	"fmt"
	"log"
	"math"
	"math/big"
	"net/http"
	"time"

	"bitbucket.org/davars/sohop/auth"
	"github.com/gorilla/mux"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

type Config struct {
	Domain          string
	Upstreams       map[string]upstreamSpec
	Github          *auth.GithubAuth
	Google          *auth.GoogleAuth
	AuthorizedOrgID int
}

type Server struct {
	Config    *Config
	CertFile  string
	CertKey   string
	HTTPAddr  string
	HTTPSAddr string

	proxy http.Handler
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func (s Server) Run() {
	var err error

	store.Options.HttpOnly = true
	store.Options.Secure = true
	store.Options.Domain = s.Config.Domain
	store.Options.MaxAge = int(sessionAge / time.Second)

	go func() {
		err = http.ListenAndServeTLS(s.HTTPSAddr, s.CertFile, s.CertKey, s.handler())
		check(err)
	}()
	go func() {
		err := http.ListenAndServe(s.HTTPAddr,
			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				r.URL.Scheme = "https"
				r.URL.Host = r.Host + s.HTTPSAddr
				http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
				return
			}))
		check(err)
	}()
	select {}
}

type upstreamSpec struct {
	URL         string
	Auth        bool
	HealthCheck string
	WebSocket   string
	Headers     http.Header
}

var (
	store       = sessions.NewCookieStore(securecookie.GenerateRandomKey(64))
	sessionName = sessionID()
	sessionAge  = 24 * time.Hour
)

func sessionID() string {
	n, err := rand.Int(rand.Reader, big.NewInt(math.MaxInt64))
	if err != nil {
		log.Fatal(err)
	}
	id := fmt.Sprintf("_s%d", n)
	return id
}

func (c *Config) authorizer() auth.Authorizer {
	if c.Github != nil && c.Google != nil {
		log.Fatal("can only use one authorizer; please configure either Google or Github authorization")
	}
	if c.Github == nil && c.Google == nil {
		log.Fatal("must define an authorizer; please configure either Google or Github authorization")
	}
	if c.Github != nil {
		return c.Github
	}
	return c.Google
}

func (s Server) handler() http.Handler {
	router := mux.NewRouter()
	router.NotFoundHandler = http.HandlerFunc(notFound)

	conf := s.Config
	oauthRouter := router.Host(fmt.Sprintf("oauth.%s", conf.Domain)).Subrouter()
	oauthRouter.Path("/authorized").Handler(auth.Handler(store, sessionName, conf.authorizer()))
	authenticating := auth.Middleware(store, sessionName, conf.authorizer())

	// TODO: switch to JWT so that this isn't necessary
	oauthRouter.Path("/session").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		session, _ := store.Get(r, sessionName)
		fmt.Fprintf(w, "%v", session.Values)
	})

	healthRouter := router.Host(fmt.Sprintf("health.%s", conf.Domain)).Subrouter()
	healthRouter.Path("/check").Handler(s.HealthHandler())

	proxyRouter := router.Host(fmt.Sprintf("{subdomain:[a-z]+}.%s", conf.Domain)).Subrouter()
	proxy := s.ProxyHandler()
	proxyRouter.MatcherFunc(requiresAuth(conf)).Handler(authenticating(proxy))
	proxyRouter.PathPrefix("/").Handler(proxy)

	return logging(router)
}
<commit_msg>Allowing subdomains to have numbers and dashes<commit_after>package sohop

import (
	"crypto/rand"
	"fmt"
	"log"
	"math"
	"math/big"
	"net/http"
	"time"

	"bitbucket.org/davars/sohop/auth"
	"github.com/gorilla/mux"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

type Config struct {
	Domain          string
	Upstreams       map[string]upstreamSpec
	Github          *auth.GithubAuth
	Google          *auth.GoogleAuth
	AuthorizedOrgID int
}

type Server struct {
	Config    *Config
	CertFile  string
	CertKey   string
	HTTPAddr  string
	HTTPSAddr string

	proxy http.Handler
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func (s Server) Run() {
	var err error

	store.Options.HttpOnly = true
	store.Options.Secure = true
	store.Options.Domain = s.Config.Domain
	store.Options.MaxAge = int(sessionAge / time.Second)

	go func() {
		err = http.ListenAndServeTLS(s.HTTPSAddr, s.CertFile, s.CertKey, s.handler())
		check(err)
	}()
	go func() {
		err := http.ListenAndServe(s.HTTPAddr,
			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				r.URL.Scheme = "https"
				r.URL.Host = r.Host + s.HTTPSAddr
				http.Redirect(w, r, r.URL.String(), http.StatusMovedPermanently)
				return
			}))
		check(err)
	}()
	select {}
}

type upstreamSpec struct {
	URL         string
	Auth        bool
	HealthCheck string
	WebSocket   string
	Headers     http.Header
}

var (
	store       = sessions.NewCookieStore(securecookie.GenerateRandomKey(64))
	sessionName = sessionID()
	sessionAge  = 24 * time.Hour
)

func sessionID() string {
	n, err := rand.Int(rand.Reader, big.NewInt(math.MaxInt64))
	if err != nil {
		log.Fatal(err)
	}
	id := fmt.Sprintf("_s%d", n)
	return id
}

func (c *Config) authorizer() auth.Authorizer {
	if c.Github != nil && c.Google != nil {
		log.Fatal("can only use one authorizer; please configure either Google or Github authorization")
	}
	if c.Github == nil && c.Google == nil {
		log.Fatal("must define an authorizer; please configure either Google or Github authorization")
	}
	if c.Github != nil {
		return c.Github
	}
	return c.Google
}

func (s Server) handler() http.Handler {
	router := mux.NewRouter()
	router.NotFoundHandler = http.HandlerFunc(notFound)

	conf := s.Config
	oauthRouter := router.Host(fmt.Sprintf("oauth.%s", conf.Domain)).Subrouter()
	oauthRouter.Path("/authorized").Handler(auth.Handler(store, sessionName, conf.authorizer()))
	authenticating := auth.Middleware(store, sessionName, conf.authorizer())

	// TODO: switch to JWT so that this isn't necessary
	oauthRouter.Path("/session").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		session, _ := store.Get(r, sessionName)
		fmt.Fprintf(w, "%v", session.Values)
	})

	healthRouter := router.Host(fmt.Sprintf("health.%s", conf.Domain)).Subrouter()
	healthRouter.Path("/check").Handler(s.HealthHandler())

	proxyRouter := router.Host(fmt.Sprintf("{subdomain:[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?}.%s", conf.Domain)).Subrouter()
	proxy := s.ProxyHandler()
	proxyRouter.MatcherFunc(requiresAuth(conf)).Handler(authenticating(proxy))
	proxyRouter.PathPrefix("/").Handler(proxy)

	return logging(router)
}
<|endoftext|>


<commit_before>// Package solve provides algorithms like A*, IDA* and Depth-first
package solve

import (
	"math"
)

// Context can be used to interact with the solver and to maintain a custom context
// during the search.
type Context struct {
	Custom interface{}
}

// The State representing a state in the search tree
//
// An implementation of this interface represents the problem. It tells the algorithm how
// to get from one state to another, how much it costs to reach the state etc.
type State interface {
	// The costs to reach this state
	Cost(ctx Context) float64

	// Returns true if this is a goal state
	IsGoal(ctx Context) bool

	// Expands this state in zero or more child states
	Expand(ctx Context) []State

	// Estimated costs to reach a goal. Use 0 for no heuristic. Most algorithms will
	// find the optimal solution if the heuristic is admissible, meaning it will never
	// over-estimate the costs to reach a goal
	Heuristic(ctx Context) float64
}

// Result of the search
type Result struct {
	// The list of states leading from the root state to the goal state. If no solution
	// is found this list will be empty
	Solution []State

	// Number of nodes visited (dequeued) by the algorithm
	Visited int

	// Number of nodes expanded (enqueued) by the algorithm
	Expanded int
}

type node struct {
	parent *node
	state  State
	value  float64
}

type result struct {
	node     *node
	contour  float64
	visited  int
	expanded int

	next *func() result
}

func generalSearch(queue strategy, visited int, expanded int, constr iconstraint, limit float64, context Context) result {
	contour := math.Inf(1)
	for {
		n := queue.Take()
		if n == nil {
			return result{nil, contour, visited, expanded, nil}
		}
		visited++
		if constr.onVisit(n) {
			continue
		}
		if n.state.IsGoal(context) {
			next := func() result {
				return generalSearch(queue, visited, expanded, constr, limit, context)
			}
			return result{n, contour, visited, expanded, &next}
		}
		for _, child := range n.state.Expand(context) {
			childNode := &node{n, child, math.Max(n.value, child.Cost(context)+child.Heuristic(context))}
			if constr.onExpand(childNode) {
				continue
			}
			if childNode.value > limit {
				contour = math.Min(contour, childNode.value)
				continue
			}
			queue.Add(childNode)
			expanded++
		}
	}
}

func idaStar(rootState State, constraint iconstraint, limit float64, context Context) result {
	visited := 0
	expanded := 0
	contour := 0.0
	for true {
		s := depthFirst()
		s.Add(&node{nil, rootState, rootState.Cost(context) + rootState.Heuristic(context)})
		constraint.reset()
		lastResult := generalSearch(s, visited, expanded, constraint, contour, context)
		lastResult.next = nil
		if lastResult.node != nil || lastResult.contour > limit || math.IsInf(lastResult.contour, 1) || math.IsNaN(lastResult.contour) {
			return lastResult
		}
		visited = lastResult.visited
		expanded = lastResult.expanded
		contour = lastResult.contour
		//fmt.Printf("contour: %v, visited: %v\n", contour, visited)
	}
	panic("Shouldn't be reached")
}

func toSlice(node *node) []State {
	if node == nil {
		return make([]State, 0)
	}
	return append(toSlice(node.parent), node.state)
}

func toResult(r *result) Result {
	return Result{toSlice(r.node), r.visited, r.expanded}
}

type solver struct {
	rootState  State
	algorithm  Algorithm
	constraint Constraint
	limit      float64
	context    interface{}

	started bool
	result  *result
}

func solve(ss *solver) Result {
	if ss.started {
		if ss.result.next == nil { // no more possible solutions
			return Result{[]State{}, ss.result.visited, ss.result.expanded}
		}
		nextResult := (*ss.result.next)()
		ss.result = &nextResult
		return toResult(ss.result)
	}
	ss.started = true
	context := Context{ss.context}
	constraint := ss.constraint.(iconstraint)
	if ss.algorithm == IDAstar {
		nextResult := idaStar(ss.rootState, constraint, ss.limit, context)
		ss.result = &nextResult
		return toResult(ss.result)
	}
	var s strategy
	switch ss.algorithm {
	case Astar:
		s = aStar()
	case DepthFirst:
		s = depthFirst()
	case BreadthFirst:
		s = breadthFirst()
	}
	s.Add(&node{nil, ss.rootState, ss.rootState.Cost(context) + ss.rootState.Heuristic(context)})

	constraint.reset()
	nextResult := generalSearch(s, 0, 0, constraint, ss.limit, context)
	ss.result = &nextResult
	return toResult(ss.result)
}

// Solver to solve the problem.
type Solver interface {
	// The algorithm to use, defaults to IDAstar
	Algorithm(algorithm Algorithm) Solver

	// The constraint to use, defaults to NONE
	Constraint(constraint Constraint) Solver

	// The limit to use. The problem will not be exanded beyond this limit. Defaults
	// to math.Inf(1).
	Limit(limit float64) Solver

	// Custom context which is passed to the methods of the state. Can contain for example precalculated data that
	// is used to speed up calculations. Be careful with state in the context though.
	Context(context interface{}) Solver

	// Solves the problem returning the result
	Solve() Result
}

func (s *solver) Algorithm(algorithm Algorithm) Solver {
	s.algorithm = algorithm
	return s
}

func (s *solver) Constraint(constraint Constraint) Solver {
	s.constraint = constraint
	return s
}

func (s *solver) Limit(limit float64) Solver {
	s.limit = limit
	return s
}

func (s *solver) Context(context interface{}) Solver {
	s.context = context
	return s
}

func (s *solver) Solve() Result {
	return solve(s)
}

// NewSolver creates a new solver
func NewSolver(rootState State) Solver {
	return &solver{rootState, Astar, NoConstraint(), math.Inf(1), nil, false, nil}
}
<commit_msg>Working on IDA* supporting multiple solutions<commit_after>// Package solve provides algorithms like A*, IDA* and Depth-first
package solve

import (
	"math"
)

// Context can be used to interact with the solver and to maintain a custom context
// during the search.
type Context struct {
	Custom interface{}
}

// The State representing a state in the search tree
//
// An implementation of this interface represents the problem. It tells the algorithm how
// to get from one state to another, how much it costs to reach the state etc.
type State interface {
	// The costs to reach this state
	Cost(ctx Context) float64

	// Returns true if this is a goal state
	IsGoal(ctx Context) bool

	// Expands this state in zero or more child states
	Expand(ctx Context) []State

	// Estimated costs to reach a goal. Use 0 for no heuristic. Most algorithms will
	// find the optimal solution if the heuristic is admissible, meaning it will never
	// over-estimate the costs to reach a goal
	Heuristic(ctx Context) float64
}

// Result of the search
type Result struct {
	// The list of states leading from the root state to the goal state. If no solution
	// is found this list will be empty
	Solution []State

	// Number of nodes visited (dequeued) by the algorithm
	Visited int

	// Number of nodes expanded (enqueued) by the algorithm
	Expanded int
}

type node struct {
	parent *node
	state  State
	value  float64
}

type result struct {
	node     *node
	contour  float64
	visited  int
	expanded int

	next *func() result
}

func generalSearch(queue strategy, visited int, expanded int, constr iconstraint, ubound float64, limit float64, contour float64, context Context) result {
	for {
		n := queue.Take()
		if n == nil {
			return result{nil, contour, visited, expanded, nil}
		}
		visited++
		if constr.onVisit(n) {
			continue
		}
		if n.state.IsGoal(context) && n.value > ubound { // IDA* regenerates previous goal nodes, we skip those using the under bound
			next := func() result {
				return generalSearch(queue, visited, expanded, constr, ubound, limit, contour, context)
			}
			return result{n, contour, visited, expanded, &next}
		}
		for _, child := range n.state.Expand(context) {
			childNode := &node{n, child, math.Max(n.value, child.Cost(context)+child.Heuristic(context))}
			if constr.onExpand(childNode) {
				continue
			}
			if childNode.value > limit {
				contour = math.Min(contour, childNode.value)
				continue
			}
			queue.Add(childNode)
			expanded++
		}
	}
}

func idaStar(rootState State, constraint iconstraint, limit float64, context Context) result {
	visited := 0
	expanded := 0
	contour := 0.0
	ubound := -1.0
	for true {
		s := depthFirst()
		s.Add(&node{nil, rootState, rootState.Cost(context) + rootState.Heuristic(context)})
		constraint.reset()
		lastResult := generalSearch(s, visited, expanded, constraint, ubound, contour, math.Inf(1), context)
		lastResult.next = nil
		if lastResult.node != nil || lastResult.contour > limit || math.IsInf(lastResult.contour, 1) || math.IsNaN(lastResult.contour) {
			return lastResult
		}
		ubound = contour
		visited = lastResult.visited
		expanded = lastResult.expanded
		contour = lastResult.contour
		//fmt.Printf("contour: %v, visited: %v\n", contour, visited)
	}
	panic("Shouldn't be reached")
}

func toSlice(node *node) []State {
	if node == nil {
		return make([]State, 0)
	}
	return append(toSlice(node.parent), node.state)
}

func toResult(r *result) Result {
	return Result{toSlice(r.node), r.visited, r.expanded}
}

type solver struct {
	rootState  State
	algorithm  Algorithm
	constraint Constraint
	limit      float64
	context    interface{}

	started bool
	result  *result
}

func solve(ss *solver) Result {
	if ss.started {
		if ss.result.next == nil { // no more possible solutions
			return Result{[]State{}, ss.result.visited, ss.result.expanded}
		}
		nextResult := (*ss.result.next)()
		ss.result = &nextResult
		return toResult(ss.result)
	}
	ss.started = true
	context := Context{ss.context}
	constraint := ss.constraint.(iconstraint)
	if ss.algorithm == IDAstar {
		nextResult := idaStar(ss.rootState, constraint, ss.limit, context)
		ss.result = &nextResult
		return toResult(ss.result)
	}
	var s strategy
	switch ss.algorithm {
	case Astar:
		s = aStar()
	case DepthFirst:
		s = depthFirst()
	case BreadthFirst:
		s = breadthFirst()
	}
	s.Add(&node{nil, ss.rootState, ss.rootState.Cost(context) + ss.rootState.Heuristic(context)})

	constraint.reset()
	nextResult := generalSearch(s, 0, 0, constraint, -1.0, ss.limit, math.Inf(1), context)
	ss.result = &nextResult
	return toResult(ss.result)
}

// Solver to solve the problem.
type Solver interface {
	// The algorithm to use, defaults to IDAstar
	Algorithm(algorithm Algorithm) Solver

	// The constraint to use, defaults to NONE
	Constraint(constraint Constraint) Solver

	// The limit to use. The problem will not be exanded beyond this limit. Defaults
	// to math.Inf(1).
	Limit(limit float64) Solver

	// Custom context which is passed to the methods of the state. Can contain for example precalculated data that
	// is used to speed up calculations. Be careful with state in the context though.
	Context(context interface{}) Solver

	// Solves the problem returning the result
	Solve() Result
}

func (s *solver) Algorithm(algorithm Algorithm) Solver {
	s.algorithm = algorithm
	return s
}

func (s *solver) Constraint(constraint Constraint) Solver {
	s.constraint = constraint
	return s
}

func (s *solver) Limit(limit float64) Solver {
	s.limit = limit
	return s
}

func (s *solver) Context(context interface{}) Solver {
	s.context = context
	return s
}

func (s *solver) Solve() Result {
	return solve(s)
}

// NewSolver creates a new solver
func NewSolver(rootState State) Solver {
	return &solver{rootState, Astar, NoConstraint(), math.Inf(1), nil, false, nil}
}
<|endoftext|>


<commit_before>package main

import (
    "bytes"
    "errors"
    "net"
    "reflect"

    log "github.com/sirupsen/logrus"
    "github.com/ugorji/go/codec"

    "github.com/nttdots/go-dots/dots_common"
    "github.com/nttdots/go-dots/dots_common/messages"
    "github.com/nttdots/go-dots/dots_server/controllers"
    "github.com/nttdots/go-dots/dots_server/models"
    "github.com/nttdots/go-dots/libcoap"
)

func unmarshalCbor(pdu *libcoap.Pdu, typ reflect.Type) (interface{}, error) {
    if len(pdu.Data) == 0 {
        return nil, nil
    }

    m := reflect.New(typ).Interface()
    reader := bytes.NewReader(pdu.Data)

    d := codec.NewDecoder(reader, dots_common.NewCborHandle())
    err := d.Decode(m)

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

func marshalCbor(msg interface{}) ([]byte, error) {
    var buf []byte
    e := codec.NewEncoderBytes(&buf, dots_common.NewCborHandle())

    err := e.Encode(msg)
    if err != nil {
        return nil, err
    }
    return buf, nil
}

func createResource(ctx *libcoap.Context, path string, typ reflect.Type, controller controllers.ControllerInterface) *libcoap.Resource {

    resource := libcoap.ResourceInit(&path, 0)
    log.Debugf("listen.go: createResource, path=%+v", path)

    var toMethodHandler = func(method controllers.ServiceMethod) libcoap.MethodHandler {
        return func(context  *libcoap.Context,
                    resource *libcoap.Resource,
                    session  *libcoap.Session,
                    request  *libcoap.Pdu,
                    token    *[]byte,
                    query    *string,
                    response *libcoap.Pdu) {

            log.WithField("MessageID", request.MessageID).Info("Incoming Request")

            response.MessageID = request.MessageID
            response.Token     = request.Token

            cn, err := session.DtlsGetPeerCommonName()
            if err != nil {
                log.WithError(err).Warn("DtlsGetPeercCommonName() failed")
                response.Code = libcoap.ResponseForbidden
                return
            }

            log.Infof("CommonName is %v", cn)

            customer, err := models.GetCustomerByCommonName(cn)
            if err != nil || customer.Id == 0 {
                log.WithError(err).Warn("Customer not found.")
                response.Code = libcoap.ResponseForbidden
                return
            }

            body, err := unmarshalCbor(request, typ)
            if err != nil {
                log.WithError(err).Error("unmarshalCbor failed.")
                response.Code = libcoap.ResponseInternalServerError
                return
            }

            req := controllers.Request {
                Code:    request.Code,
                Type:    request.Type,
                Uri:     request.Path(),
                Queries: request.Queries(),
                Body:    body,
            }
            log.Debugf("req=%+v", req)

            res, err := method(req, customer)
            if err != nil {
                log.WithError(err).Error("controller returned error")
                response.Code = libcoap.ResponseInternalServerError
                return
            }
            log.Debugf("res=%+v", res)
            payload, err := marshalCbor(res.Body)
            if err != nil {
                log.WithError(err).Error("marshalCbor failed.")
                response.Code = libcoap.ResponseInternalServerError
                return
            }

            response.Code = libcoap.Code(res.Code)
            response.Data = payload

            return
        }
    }

    resource.RegisterHandler(libcoap.RequestGet,    toMethodHandler(controller.HandleGet))
    resource.RegisterHandler(libcoap.RequestPut,    toMethodHandler(controller.HandlePut))
    resource.RegisterHandler(libcoap.RequestPost,   toMethodHandler(controller.HandlePost))
    resource.RegisterHandler(libcoap.RequestDelete, toMethodHandler(controller.HandleDelete))
    return resource
}

func addHandler(ctx *libcoap.Context, code messages.Code, controller controllers.ControllerInterface) {
    msg := messages.MessageTypes[code]
    path := "/" + msg.Path

    ctx.AddResource(createResource(ctx, path, msg.Type, controller))
}

func addPrefixHandler(ctx *libcoap.Context, code messages.Code, controller controllers.ControllerInterface) {
    msg := messages.MessageTypes[code]
    path := "/" + msg.Path

    filter := controllers.NewPrefixFilter(path, controller)
    ctx.AddResourceUnknown(createResource(ctx, "dummy for unknown", msg.Type, filter))
}

func listen(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    log.Debugf("listen.go, listen -in. address=%+v, port=%+v", address, port)
    ip := net.ParseIP(address)
    if ip == nil {
        err = errors.New("net.ParseIP() -> nil")
        return
    }

    addr, err := libcoap.AddressOf(ip, port)
    if err != nil {
        return
    }
    log.Debugf("addr=%+v", addr)

    ctx := libcoap.NewContextDtls(nil, dtlsParam)
    if ctx == nil {
        err = errors.New("libcoap.NewContextDtls() -> nil")
        return
    }

    ctx.NewEndpoint(addr, libcoap.ProtoDtls)
    return ctx, nil
}

func listenData(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    ctx, err := listen(address, port, dtlsParam)
    if err != nil {
        return
    }

    addHandler(ctx, messages.HELLO,                  &controllers.Hello{})
    addHandler(ctx, messages.CREATE_IDENTIFIERS,     &controllers.CreateIdentifiers{})
    addHandler(ctx, messages.INSTALL_FILTERING_RULE, &controllers.InstallFilteringRule{})

    return ctx, nil
}

func listenSignal(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    ctx, err := listen(address, port, dtlsParam)
    if err != nil {
        return
    }

    addHandler(ctx, messages.HELLO,                 &controllers.Hello{})
    addHandler(ctx, messages.SESSION_CONFIGURATION, &controllers.SessionConfiguration{})

    addPrefixHandler(ctx, messages.MITIGATION_REQUEST, &controllers.MitigationRequest{})

    return ctx, nil
}
<commit_msg>Add ContentType CBOR on response of CoAP<commit_after>package main

import (
    "bytes"
    "errors"
    "net"
    "reflect"

    log "github.com/sirupsen/logrus"
    "github.com/ugorji/go/codec"

    "github.com/nttdots/go-dots/dots_common"
    "github.com/nttdots/go-dots/dots_common/messages"
    "github.com/nttdots/go-dots/dots_server/controllers"
    "github.com/nttdots/go-dots/dots_server/models"
    "github.com/nttdots/go-dots/libcoap"
)

func unmarshalCbor(pdu *libcoap.Pdu, typ reflect.Type) (interface{}, error) {
    if len(pdu.Data) == 0 {
        return nil, nil
    }

    m := reflect.New(typ).Interface()
    reader := bytes.NewReader(pdu.Data)

    d := codec.NewDecoder(reader, dots_common.NewCborHandle())
    err := d.Decode(m)

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

func marshalCbor(msg interface{}) ([]byte, error) {
    var buf []byte
    e := codec.NewEncoderBytes(&buf, dots_common.NewCborHandle())

    err := e.Encode(msg)
    if err != nil {
        return nil, err
    }
    return buf, nil
}

func createResource(ctx *libcoap.Context, path string, typ reflect.Type, controller controllers.ControllerInterface) *libcoap.Resource {

    resource := libcoap.ResourceInit(&path, 0)
    log.Debugf("listen.go: createResource, path=%+v", path)

    var toMethodHandler = func(method controllers.ServiceMethod) libcoap.MethodHandler {
        return func(context  *libcoap.Context,
                    resource *libcoap.Resource,
                    session  *libcoap.Session,
                    request  *libcoap.Pdu,
                    token    *[]byte,
                    query    *string,
                    response *libcoap.Pdu) {

            log.WithField("MessageID", request.MessageID).Info("Incoming Request")

            response.MessageID = request.MessageID
            response.Token     = request.Token

            cn, err := session.DtlsGetPeerCommonName()
            if err != nil {
                log.WithError(err).Warn("DtlsGetPeercCommonName() failed")
                response.Code = libcoap.ResponseForbidden
                return
            }

            log.Infof("CommonName is %v", cn)

            customer, err := models.GetCustomerByCommonName(cn)
            if err != nil || customer.Id == 0 {
                log.WithError(err).Warn("Customer not found.")
                response.Code = libcoap.ResponseForbidden
                return
            }

            body, err := unmarshalCbor(request, typ)
            if err != nil {
                log.WithError(err).Error("unmarshalCbor failed.")
                response.Code = libcoap.ResponseInternalServerError
                return
            }

            req := controllers.Request {
                Code:    request.Code,
                Type:    request.Type,
                Uri:     request.Path(),
                Queries: request.Queries(),
                Body:    body,
            }
            log.Debugf("req=%+v", req)

            res, err := method(req, customer)
            if err != nil {
                log.WithError(err).Error("controller returned error")
                response.Code = libcoap.ResponseInternalServerError
                return
            }

            log.Debugf("res=%+v", res)
            payload, err := marshalCbor(res.Body)
            if err != nil {
                log.WithError(err).Error("marshalCbor failed.")
                response.Code = libcoap.ResponseInternalServerError
                return
            }

            response.Code = libcoap.Code(res.Code)
            response.Data = payload
            // add content type cbor
            response.Options = append(response.Options, libcoap.OptionContentType.Uint16(60))

            return
        }
    }

    resource.RegisterHandler(libcoap.RequestGet,    toMethodHandler(controller.HandleGet))
    resource.RegisterHandler(libcoap.RequestPut,    toMethodHandler(controller.HandlePut))
    resource.RegisterHandler(libcoap.RequestPost,   toMethodHandler(controller.HandlePost))
    resource.RegisterHandler(libcoap.RequestDelete, toMethodHandler(controller.HandleDelete))
    return resource
}

func addHandler(ctx *libcoap.Context, code messages.Code, controller controllers.ControllerInterface) {
    msg := messages.MessageTypes[code]
    path := "/" + msg.Path

    ctx.AddResource(createResource(ctx, path, msg.Type, controller))
}

func addPrefixHandler(ctx *libcoap.Context, code messages.Code, controller controllers.ControllerInterface) {
    msg := messages.MessageTypes[code]
    path := "/" + msg.Path

    filter := controllers.NewPrefixFilter(path, controller)
    ctx.AddResourceUnknown(createResource(ctx, "dummy for unknown", msg.Type, filter))
}

func listen(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    log.Debugf("listen.go, listen -in. address=%+v, port=%+v", address, port)
    ip := net.ParseIP(address)
    if ip == nil {
        err = errors.New("net.ParseIP() -> nil")
        return
    }

    addr, err := libcoap.AddressOf(ip, port)
    if err != nil {
        return
    }
    log.Debugf("addr=%+v", addr)

    ctx := libcoap.NewContextDtls(nil, dtlsParam)
    if ctx == nil {
        err = errors.New("libcoap.NewContextDtls() -> nil")
        return
    }

    ctx.NewEndpoint(addr, libcoap.ProtoDtls)
    return ctx, nil
}

func listenData(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    ctx, err := listen(address, port, dtlsParam)
    if err != nil {
        return
    }

    addHandler(ctx, messages.HELLO,                  &controllers.Hello{})
    addHandler(ctx, messages.CREATE_IDENTIFIERS,     &controllers.CreateIdentifiers{})
    addHandler(ctx, messages.INSTALL_FILTERING_RULE, &controllers.InstallFilteringRule{})

    return ctx, nil
}

func listenSignal(address string, port uint16, dtlsParam *libcoap.DtlsParam) (_ *libcoap.Context, err error) {
    ctx, err := listen(address, port, dtlsParam)
    if err != nil {
        return
    }

    addHandler(ctx, messages.HELLO,                 &controllers.Hello{})
    addHandler(ctx, messages.SESSION_CONFIGURATION, &controllers.SessionConfiguration{})

    addPrefixHandler(ctx, messages.MITIGATION_REQUEST, &controllers.MitigationRequest{})

    return ctx, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

type UserSuite struct {
	ConnSuite
}

var _ = gc.Suite(&UserSuite{})

func (s *UserSuite) TestAddUserInvalidNames(c *gc.C) {
	for _, name := range []string{
		"foo-bar",
		"",
		"0foo",
	} {
		u, err := s.State.AddUser(name, "password")
		c.Assert(err, gc.ErrorMatches, `invalid user name "`+name+`"`)
		c.Assert(u, gc.IsNil)
	}
}

func (s *UserSuite) TestAddUser(c *gc.C) {
	u, err := s.State.AddUser("a", "b")
	c.Check(u, gc.NotNil)
	c.Assert(err, gc.IsNil)

	c.Assert(u.Name(), gc.Equals, "a")
	c.Assert(u.PasswordValid("b"), jc.IsTrue)

	u1, err := s.State.User("a")
	c.Check(u1, gc.NotNil)
	c.Assert(err, gc.IsNil)

	c.Assert(u1.Name(), gc.Equals, "a")
	c.Assert(u1.PasswordValid("b"), jc.IsTrue)
}

func (s *UserSuite) TestCheckUserExists(c *gc.C) {
	u, err := s.State.AddUser("a", "b")
	c.Check(u, gc.NotNil)
	c.Assert(err, gc.IsNil)
	e, err := state.CheckUserExists(s.State, "a")
	c.Assert(err, gc.IsNil)
	c.Assert(e, gc.Equals, true)
	e, err = state.CheckUserExists(s.State, "notAUser")
	c.Assert(err, gc.IsNil)
	c.Assert(e, gc.Equals, false)
}

func (s *UserSuite) TestSetPassword(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	testSetPassword(c, func() (state.Authenticator, error) {
		return s.State.User(u.Name())
	})
}

func (s *UserSuite) TestAddUserSetsSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "a-password")
	c.Assert(err, gc.IsNil)
	salt, hash := state.GetUserPasswordSaltAndHash(u)
	c.Check(hash, gc.Not(gc.Equals), "")
	c.Check(salt, gc.Not(gc.Equals), "")
	c.Check(utils.UserPasswordHash("a-password", salt), gc.Equals, hash)
	c.Check(u.PasswordValid("a-password"), jc.IsTrue)
}

func (s *UserSuite) TestSetPasswordChangesSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "a-password")
	c.Assert(err, gc.IsNil)
	origSalt, origHash := state.GetUserPasswordSaltAndHash(u)
	c.Check(origSalt, gc.Not(gc.Equals), "")
	// Even though the password is the same, we take this opportunity to
	// update the salt
	u.SetPassword("a-password")
	newSalt, newHash := state.GetUserPasswordSaltAndHash(u)
	c.Check(newSalt, gc.Not(gc.Equals), "")
	c.Check(newSalt, gc.Not(gc.Equals), origSalt)
	c.Check(newHash, gc.Not(gc.Equals), origHash)
	c.Check(u.PasswordValid("a-password"), jc.IsTrue)
}

func (s *UserSuite) TestSetPasswordHash(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	err = u.SetPasswordHash(utils.UserPasswordHash("foo", utils.CompatSalt), utils.CompatSalt)
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	c.Assert(u.PasswordValid("bar"), jc.IsFalse)

	// User passwords should *not* use the fast PasswordHash function
	hash := utils.AgentPasswordHash("foo-12345678901234567890")
	c.Assert(err, gc.IsNil)
	err = u.SetPasswordHash(hash, "")
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo-12345678901234567890"), jc.IsFalse)
}

func (s *UserSuite) TestSetPasswordHashWithSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	err = u.SetPasswordHash(utils.UserPasswordHash("foo", "salted"), "salted")
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	salt, hash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(salt, gc.Equals, "salted")
	c.Assert(hash, gc.Not(gc.Equals), utils.UserPasswordHash("foo", utils.CompatSalt))
}

func (s *UserSuite) TestPasswordValidUpdatesSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	compatHash := utils.UserPasswordHash("foo", utils.CompatSalt)
	err = u.SetPasswordHash(compatHash, "")
	c.Assert(err, gc.IsNil)
	beforeSalt, beforeHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(beforeSalt, gc.Equals, "")
	c.Assert(beforeHash, gc.Equals, compatHash)
	c.Assert(u.PasswordValid("bar"), jc.IsFalse)
	// A bad password doesn't trigger a rewrite
	afterBadSalt, afterBadHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(afterBadSalt, gc.Equals, "")
	c.Assert(afterBadHash, gc.Equals, compatHash)
	// When we get a valid check, we then add a salt and rewrite the hash
	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	afterSalt, afterHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(afterSalt, gc.Not(gc.Equals), "")
	c.Assert(afterHash, gc.Not(gc.Equals), compatHash)
	c.Assert(afterHash, gc.Equals, utils.UserPasswordHash("foo", afterSalt))
	// running PasswordValid again doesn't trigger another rewrite
	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	lastSalt, lastHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(lastSalt, gc.Equals, afterSalt)
	c.Assert(lastHash, gc.Equals, afterHash)
}

func (s *UserSuite) TestName(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	c.Assert(u.Name(), gc.Equals, "someuser")
	c.Assert(u.Tag(), gc.Equals, "user-someuser")
}

func (s *UserSuite) TestDeactivate(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)
	c.Assert(u.IsDeactivated(), gc.Equals, false)

	err = u.Deactivate()
	c.Assert(err, gc.IsNil)
	c.Assert(u.IsDeactivated(), gc.Equals, true)
	c.Assert(u.PasswordValid(""), gc.Equals, false)

}

func (s *UserSuite) TestCantDeactivateAdminUser(c *gc.C) {
	u, err := s.State.User(state.AdminUser)
	c.Assert(err, gc.IsNil)
	err = u.Deactivate()
	c.Assert(err, gc.ErrorMatches, "Can't deactivate admin user")
}
<commit_msg>Update/add tests<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"regexp"

	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

type UserSuite struct {
	ConnSuite
}

var _ = gc.Suite(&UserSuite{})

func (s *UserSuite) TestAddUserInvalidNames(c *gc.C) {
	for _, name := range []string{
		"bar@ram.u",
		"#1foo",
		"&bar",
		"%bar",
		"foo()",
		"'foo'",
		"[bar]",
		"foo?",
		"foo=bar",
		"bar*",
		"bar^",
		"bar!",
		"bar_foo",
		"bar+foo",
		"bar{}",
		"foo bar",
		"",
	} {
		u, err := s.State.AddUser(name, "password")
		c.Assert(err, gc.ErrorMatches, `invalid user name "`+regexp.QuoteMeta(name)+`"`)
		c.Assert(u, gc.IsNil)
	}
}

func (s *UserSuite) TestAddUserValidName(c *gc.C) {
	name := "007-Foo.Bar"
	u, err := s.State.AddUser(name, "password")
	c.Check(u, gc.NotNil)
	c.Assert(err, gc.IsNil)
	c.Assert(u.Name(), gc.Equals, name)
}

func (s *UserSuite) TestAddUser(c *gc.C) {
	u, err := s.State.AddUser("a", "b")
	c.Check(u, gc.NotNil)
	c.Assert(err, gc.IsNil)

	c.Assert(u.Name(), gc.Equals, "a")
	c.Assert(u.PasswordValid("b"), jc.IsTrue)

	u1, err := s.State.User("a")
	c.Check(u1, gc.NotNil)
	c.Assert(err, gc.IsNil)

	c.Assert(u1.Name(), gc.Equals, "a")
	c.Assert(u1.PasswordValid("b"), jc.IsTrue)
}

func (s *UserSuite) TestCheckUserExists(c *gc.C) {
	u, err := s.State.AddUser("a", "b")
	c.Check(u, gc.NotNil)
	c.Assert(err, gc.IsNil)
	e, err := state.CheckUserExists(s.State, "a")
	c.Assert(err, gc.IsNil)
	c.Assert(e, gc.Equals, true)
	e, err = state.CheckUserExists(s.State, "notAUser")
	c.Assert(err, gc.IsNil)
	c.Assert(e, gc.Equals, false)
}

func (s *UserSuite) TestSetPassword(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	testSetPassword(c, func() (state.Authenticator, error) {
		return s.State.User(u.Name())
	})
}

func (s *UserSuite) TestAddUserSetsSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "a-password")
	c.Assert(err, gc.IsNil)
	salt, hash := state.GetUserPasswordSaltAndHash(u)
	c.Check(hash, gc.Not(gc.Equals), "")
	c.Check(salt, gc.Not(gc.Equals), "")
	c.Check(utils.UserPasswordHash("a-password", salt), gc.Equals, hash)
	c.Check(u.PasswordValid("a-password"), jc.IsTrue)
}

func (s *UserSuite) TestSetPasswordChangesSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "a-password")
	c.Assert(err, gc.IsNil)
	origSalt, origHash := state.GetUserPasswordSaltAndHash(u)
	c.Check(origSalt, gc.Not(gc.Equals), "")
	// Even though the password is the same, we take this opportunity to
	// update the salt
	u.SetPassword("a-password")
	newSalt, newHash := state.GetUserPasswordSaltAndHash(u)
	c.Check(newSalt, gc.Not(gc.Equals), "")
	c.Check(newSalt, gc.Not(gc.Equals), origSalt)
	c.Check(newHash, gc.Not(gc.Equals), origHash)
	c.Check(u.PasswordValid("a-password"), jc.IsTrue)
}

func (s *UserSuite) TestSetPasswordHash(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	err = u.SetPasswordHash(utils.UserPasswordHash("foo", utils.CompatSalt), utils.CompatSalt)
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	c.Assert(u.PasswordValid("bar"), jc.IsFalse)

	// User passwords should *not* use the fast PasswordHash function
	hash := utils.AgentPasswordHash("foo-12345678901234567890")
	c.Assert(err, gc.IsNil)
	err = u.SetPasswordHash(hash, "")
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo-12345678901234567890"), jc.IsFalse)
}

func (s *UserSuite) TestSetPasswordHashWithSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	err = u.SetPasswordHash(utils.UserPasswordHash("foo", "salted"), "salted")
	c.Assert(err, gc.IsNil)

	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	salt, hash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(salt, gc.Equals, "salted")
	c.Assert(hash, gc.Not(gc.Equals), utils.UserPasswordHash("foo", utils.CompatSalt))
}

func (s *UserSuite) TestPasswordValidUpdatesSalt(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	compatHash := utils.UserPasswordHash("foo", utils.CompatSalt)
	err = u.SetPasswordHash(compatHash, "")
	c.Assert(err, gc.IsNil)
	beforeSalt, beforeHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(beforeSalt, gc.Equals, "")
	c.Assert(beforeHash, gc.Equals, compatHash)
	c.Assert(u.PasswordValid("bar"), jc.IsFalse)
	// A bad password doesn't trigger a rewrite
	afterBadSalt, afterBadHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(afterBadSalt, gc.Equals, "")
	c.Assert(afterBadHash, gc.Equals, compatHash)
	// When we get a valid check, we then add a salt and rewrite the hash
	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	afterSalt, afterHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(afterSalt, gc.Not(gc.Equals), "")
	c.Assert(afterHash, gc.Not(gc.Equals), compatHash)
	c.Assert(afterHash, gc.Equals, utils.UserPasswordHash("foo", afterSalt))
	// running PasswordValid again doesn't trigger another rewrite
	c.Assert(u.PasswordValid("foo"), jc.IsTrue)
	lastSalt, lastHash := state.GetUserPasswordSaltAndHash(u)
	c.Assert(lastSalt, gc.Equals, afterSalt)
	c.Assert(lastHash, gc.Equals, afterHash)
}

func (s *UserSuite) TestName(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)

	c.Assert(u.Name(), gc.Equals, "someuser")
	c.Assert(u.Tag(), gc.Equals, "user-someuser")
}

func (s *UserSuite) TestDeactivate(c *gc.C) {
	u, err := s.State.AddUser("someuser", "password")
	c.Assert(err, gc.IsNil)
	c.Assert(u.IsDeactivated(), gc.Equals, false)

	err = u.Deactivate()
	c.Assert(err, gc.IsNil)
	c.Assert(u.IsDeactivated(), gc.Equals, true)
	c.Assert(u.PasswordValid(""), gc.Equals, false)

}

func (s *UserSuite) TestCantDeactivateAdminUser(c *gc.C) {
	u, err := s.State.User(state.AdminUser)
	c.Assert(err, gc.IsNil)
	err = u.Deactivate()
	c.Assert(err, gc.ErrorMatches, "Can't deactivate admin user")
}
<|endoftext|>


<commit_before>package webdriver

import (
	"testing"

	. "github.com/bborbe/assert"
	monitoring_check "github.com/bborbe/monitoring/check"
)

func TestImplementsCheck(t *testing.T) {
	c := New(nil,"http://www.example.com")
	var i *monitoring_check.Check
	err := AssertThat(c, Implements(i))
	if err != nil {
		t.Fatal(err)
	}
}

func TestDescription(t *testing.T) {
	c := New(nil,"http://www.example.com")
	err := AssertThat(c.Description(), Is("webdriver check on url http://www.example.com"))
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>format<commit_after>package webdriver

import (
	"testing"

	. "github.com/bborbe/assert"
	monitoring_check "github.com/bborbe/monitoring/check"
)

func TestImplementsCheck(t *testing.T) {
	c := New(nil, "http://www.example.com")
	var i *monitoring_check.Check
	err := AssertThat(c, Implements(i))
	if err != nil {
		t.Fatal(err)
	}
}

func TestDescription(t *testing.T) {
	c := New(nil, "http://www.example.com")
	err := AssertThat(c.Description(), Is("webdriver check on url http://www.example.com"))
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/eval"
	"github.com/elpinal/coco3/gate"
	"github.com/elpinal/coco3/parser"

	"github.com/elpinal/coco3/extra"
	eparser "github.com/elpinal/coco3/extra/parser"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

type CLI struct {
	In  io.Reader
	Out io.Writer
	Err io.Writer

	*config.Config

	db *sqlx.DB

	exitCh chan int
	doneCh chan struct{} // to ensure exiting just after exitCh received

	execute1 func([]byte) error
}

func (c *CLI) Run(args []string) int {
	f := flag.NewFlagSet("coco3", flag.ContinueOnError)
	f.SetOutput(c.Err)
	f.Usage = func() {
		c.Err.Write([]byte("coco3 is a shell.\n"))
		c.Err.Write([]byte("Usage:\n"))
		f.PrintDefaults()
	}

	if c.Config == nil {
		c.Config = &config.Config{}
	}
	c.Config.Init()
	flagC := f.String("c", "", "take first argument as a command to execute")
	flagE := f.Bool("extra", c.Config.Extra, "switch to extra mode")
	if err := f.Parse(args); err != nil {
		return 2
	}
	return c.run(f.Args(), flagC, flagE)
}

func (c *CLI) run(args []string, flagC *string, flagE *bool) int {
	c.exitCh = make(chan int)
	c.doneCh = make(chan struct{})
	defer close(c.doneCh)

	for _, alias := range c.Config.Alias {
		eval.DefAlias(alias[0], alias[1])
	}

	for k, v := range c.Config.Env {
		err := os.Setenv(k, v)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			return 1
		}
	}

	setpath(c.Config.Paths)

	if *flagE {
		// If -extra flag is on, enable extra mode on any command executions.
		c.execute1 = c.executeExtra
	} else {
		c.execute1 = c.execute
	}

	if len(c.Config.StartUpCommand) > 0 {
		done := make(chan struct{})
		go func() {
			if err := c.execute1(c.Config.StartUpCommand); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
			}
			close(done)
		}()
		select {
		case code := <-c.exitCh:
			return code
		case <-done:
		}
	}

	if *flagC != "" {
		go func() {
			if err := c.execute1([]byte(*flagC)); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
				return
			}
			c.exitCh <- 0
		}()
		return <-c.exitCh
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if len(args) > 0 {
		go c.runFiles(ctx, args)
		return <-c.exitCh
	}

	histRunes, err := c.getHistory(c.Config.HistFile)
	if err != nil {
		fmt.Fprintln(c.Err, err)
		return 1
	}
	g := gate.NewContext(ctx, c.Config, c.In, c.Out, c.Err, histRunes)
	go func(ctx context.Context) {
		for {
			if err := c.interact(g); err != nil {
				c.printExecError(err)
				g.Clear()
			}
			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}(ctx)
	return <-c.exitCh
}

func (c *CLI) errorf(s string) {
	fmt.Fprintf(c.Err, s)
}

func (c *CLI) errorln(s interface{}) {
	fmt.Fprintln(c.Err, s)
}

func (c *CLI) errorp(s ...interface{}) {
	fmt.Fprint(c.Err, s...)
}

func (c *CLI) getHistory(filename string) ([][]rune, error) {
	db, err := sqlx.Connect("sqlite3", filename)
	if err != nil {
		return nil, errors.Wrap(err, "connecting history file")
	}
	_, err = db.Exec(schema)
	if err != nil {
		return nil, errors.Wrap(err, "initializing history file")
	}
	var history []string
	err = db.Select(&history, "select line from command_info")
	if err != nil {
		return nil, errors.Wrap(err, "restoring history")
	}
	// TODO: Is this way proper?
	c.db = db
	return sanitizeHistory(history), nil
}

func (c *CLI) printExecError(err error) {
	if pe, ok := err.(*eparser.ParseError); ok {
		fmt.Fprintln(c.Err, pe.Verbose())
	} else {
		c.errorln(err)
	}
}

// setpath sets the PATH environment variable.
func setpath(args []string) {
	if len(args) == 0 {
		return
	}
	paths := filepath.SplitList(os.Getenv("PATH"))
	var newPaths []string
	for _, path := range paths {
		if contains(args, path) {
			continue
		}
		newPaths = append(newPaths, path)
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, string(filepath.ListSeparator)))
}

func contains(xs []string, s string) bool {
	for _, x := range xs {
		if x == s {
			return true
		}
	}
	return false
}

func sanitizeHistory(history []string) [][]rune {
	histRunes := make([][]rune, 0, len(history))
	for _, line := range history {
		if line == "" {
			continue
		}
		l := len(histRunes)
		s := []rune(line)
		if l > 0 && compareRunes(histRunes[l-1], s) {
			continue
		}
		histRunes = append(histRunes, s)
	}
	return histRunes
}

func compareRunes(r1, r2 []rune) bool {
	if len(r1) != len(r2) {
		return false
	}
	for i, r := range r1 {
		if r2[i] != r {
			return false
		}
	}
	return true
}

func (c *CLI) interact(g gate.Gate) error {
	r, end, err := c.read(g)
	if err != nil {
		return err
	}
	if end {
		c.exitCh <- 0
		<-c.doneCh
		return nil
	}
	go c.writeHistory(r)
	if err := c.execute1([]byte(string(r))); err != nil {
		return err
	}
	g.Clear()
	return nil
}

func (c *CLI) read(g gate.Gate) ([]rune, bool, error) {
	defer c.Out.Write([]byte{'\n'})
	oldState, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := terminal.Restore(0, oldState); err != nil {
			fmt.Fprintln(c.Err, err)
		}
	}()
	r, end, err := g.Read()
	if err != nil {
		return nil, false, err
	}
	return r, end, nil
}

func (c *CLI) writeHistory(r []rune) {
	startTime := time.Now()
	_, err := c.db.Exec("insert into command_info (time, line) values ($1, $2)", startTime, string(r))
	if err != nil {
		fmt.Fprintf(c.Err, "saving history: %v\n", err)
		c.exitCh <- 1
	}
}

const schema = `
create table if not exists command_info (
    time datetime,
    line text
)`

func (c *CLI) execute(b []byte) error {
	f, err := parser.ParseSrc(b)
	if err != nil {
		return err
	}
	e := eval.New(c.In, c.Out, c.Err, c.db)
	err = e.Eval(f.Lines)
	select {
	case code := <-e.ExitCh:
		c.exitCh <- code
		<-c.doneCh
	default:
	}
	return err
}

func (c *CLI) executeExtra(b []byte) error {
	cmd, err := eparser.Parse(b)
	if err != nil {
		return err
	}
	e := extra.New(extra.Option{DB: c.db})
	err = e.Eval(cmd)
	if err == nil {
		return nil
	}
	if pe, ok := err.(*eparser.ParseError); ok {
		pe.Src = string(b)
	}
	return err
}

func (c *CLI) runFiles(ctx context.Context, files []string) {
	for _, file := range files {
		b, err := ioutil.ReadFile(file)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		if err := c.execute1(b); err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}
	}
	c.exitCh <- 0
}
<commit_msg>Change errorf and errorln<commit_after>package cli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/eval"
	"github.com/elpinal/coco3/gate"
	"github.com/elpinal/coco3/parser"

	"github.com/elpinal/coco3/extra"
	eparser "github.com/elpinal/coco3/extra/parser"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

type CLI struct {
	In  io.Reader
	Out io.Writer
	Err io.Writer

	*config.Config

	db *sqlx.DB

	exitCh chan int
	doneCh chan struct{} // to ensure exiting just after exitCh received

	execute1 func([]byte) error
}

func (c *CLI) Run(args []string) int {
	f := flag.NewFlagSet("coco3", flag.ContinueOnError)
	f.SetOutput(c.Err)
	f.Usage = func() {
		c.Err.Write([]byte("coco3 is a shell.\n"))
		c.Err.Write([]byte("Usage:\n"))
		f.PrintDefaults()
	}

	if c.Config == nil {
		c.Config = &config.Config{}
	}
	c.Config.Init()
	flagC := f.String("c", "", "take first argument as a command to execute")
	flagE := f.Bool("extra", c.Config.Extra, "switch to extra mode")
	if err := f.Parse(args); err != nil {
		return 2
	}
	return c.run(f.Args(), flagC, flagE)
}

func (c *CLI) run(args []string, flagC *string, flagE *bool) int {
	c.exitCh = make(chan int)
	c.doneCh = make(chan struct{})
	defer close(c.doneCh)

	for _, alias := range c.Config.Alias {
		eval.DefAlias(alias[0], alias[1])
	}

	for k, v := range c.Config.Env {
		err := os.Setenv(k, v)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			return 1
		}
	}

	setpath(c.Config.Paths)

	if *flagE {
		// If -extra flag is on, enable extra mode on any command executions.
		c.execute1 = c.executeExtra
	} else {
		c.execute1 = c.execute
	}

	if len(c.Config.StartUpCommand) > 0 {
		done := make(chan struct{})
		go func() {
			if err := c.execute1(c.Config.StartUpCommand); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
			}
			close(done)
		}()
		select {
		case code := <-c.exitCh:
			return code
		case <-done:
		}
	}

	if *flagC != "" {
		go func() {
			if err := c.execute1([]byte(*flagC)); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
				return
			}
			c.exitCh <- 0
		}()
		return <-c.exitCh
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if len(args) > 0 {
		go c.runFiles(ctx, args)
		return <-c.exitCh
	}

	histRunes, err := c.getHistory(c.Config.HistFile)
	if err != nil {
		fmt.Fprintln(c.Err, err)
		return 1
	}
	g := gate.NewContext(ctx, c.Config, c.In, c.Out, c.Err, histRunes)
	go func(ctx context.Context) {
		for {
			if err := c.interact(g); err != nil {
				c.printExecError(err)
				g.Clear()
			}
			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}(ctx)
	return <-c.exitCh
}

func (c *CLI) errorf(s string, a ...interface{}) {
	fmt.Fprintf(c.Err, s, a...)
}

func (c *CLI) errorln(s ...interface{}) {
	fmt.Fprintln(c.Err, s...)
}

func (c *CLI) errorp(s ...interface{}) {
	fmt.Fprint(c.Err, s...)
}

func (c *CLI) getHistory(filename string) ([][]rune, error) {
	db, err := sqlx.Connect("sqlite3", filename)
	if err != nil {
		return nil, errors.Wrap(err, "connecting history file")
	}
	_, err = db.Exec(schema)
	if err != nil {
		return nil, errors.Wrap(err, "initializing history file")
	}
	var history []string
	err = db.Select(&history, "select line from command_info")
	if err != nil {
		return nil, errors.Wrap(err, "restoring history")
	}
	// TODO: Is this way proper?
	c.db = db
	return sanitizeHistory(history), nil
}

func (c *CLI) printExecError(err error) {
	if pe, ok := err.(*eparser.ParseError); ok {
		fmt.Fprintln(c.Err, pe.Verbose())
	} else {
		c.errorln(err)
	}
}

// setpath sets the PATH environment variable.
func setpath(args []string) {
	if len(args) == 0 {
		return
	}
	paths := filepath.SplitList(os.Getenv("PATH"))
	var newPaths []string
	for _, path := range paths {
		if contains(args, path) {
			continue
		}
		newPaths = append(newPaths, path)
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, string(filepath.ListSeparator)))
}

func contains(xs []string, s string) bool {
	for _, x := range xs {
		if x == s {
			return true
		}
	}
	return false
}

func sanitizeHistory(history []string) [][]rune {
	histRunes := make([][]rune, 0, len(history))
	for _, line := range history {
		if line == "" {
			continue
		}
		l := len(histRunes)
		s := []rune(line)
		if l > 0 && compareRunes(histRunes[l-1], s) {
			continue
		}
		histRunes = append(histRunes, s)
	}
	return histRunes
}

func compareRunes(r1, r2 []rune) bool {
	if len(r1) != len(r2) {
		return false
	}
	for i, r := range r1 {
		if r2[i] != r {
			return false
		}
	}
	return true
}

func (c *CLI) interact(g gate.Gate) error {
	r, end, err := c.read(g)
	if err != nil {
		return err
	}
	if end {
		c.exitCh <- 0
		<-c.doneCh
		return nil
	}
	go c.writeHistory(r)
	if err := c.execute1([]byte(string(r))); err != nil {
		return err
	}
	g.Clear()
	return nil
}

func (c *CLI) read(g gate.Gate) ([]rune, bool, error) {
	defer c.Out.Write([]byte{'\n'})
	oldState, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := terminal.Restore(0, oldState); err != nil {
			fmt.Fprintln(c.Err, err)
		}
	}()
	r, end, err := g.Read()
	if err != nil {
		return nil, false, err
	}
	return r, end, nil
}

func (c *CLI) writeHistory(r []rune) {
	startTime := time.Now()
	_, err := c.db.Exec("insert into command_info (time, line) values ($1, $2)", startTime, string(r))
	if err != nil {
		fmt.Fprintf(c.Err, "saving history: %v\n", err)
		c.exitCh <- 1
	}
}

const schema = `
create table if not exists command_info (
    time datetime,
    line text
)`

func (c *CLI) execute(b []byte) error {
	f, err := parser.ParseSrc(b)
	if err != nil {
		return err
	}
	e := eval.New(c.In, c.Out, c.Err, c.db)
	err = e.Eval(f.Lines)
	select {
	case code := <-e.ExitCh:
		c.exitCh <- code
		<-c.doneCh
	default:
	}
	return err
}

func (c *CLI) executeExtra(b []byte) error {
	cmd, err := eparser.Parse(b)
	if err != nil {
		return err
	}
	e := extra.New(extra.Option{DB: c.db})
	err = e.Eval(cmd)
	if err == nil {
		return nil
	}
	if pe, ok := err.(*eparser.ParseError); ok {
		pe.Src = string(b)
	}
	return err
}

func (c *CLI) runFiles(ctx context.Context, files []string) {
	for _, file := range files {
		b, err := ioutil.ReadFile(file)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		if err := c.execute1(b); err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}
	}
	c.exitCh <- 0
}
<|endoftext|>


<commit_before>package jujutest

import (
	"bytes"
	"io"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/trivial"
	"net/http"
)

// Tests is a gocheck suite containing tests verifying juju functionality
// against the environment with the given configuration. The
// tests are not designed to be run against a live server - the Environ
// is opened once for each test, and some potentially expensive operations
// may be executed.
type Tests struct {
	coretesting.LoggingSuite
	Config map[string]interface{}
	Env    environs.Environ
}

// Open opens an instance of the testing environment.
func (t *Tests) Open(c *C) environs.Environ {
	e, err := environs.NewFromAttrs(t.Config)
	c.Assert(err, IsNil, Commentf("opening environ %#v", t.Config))
	c.Assert(e, NotNil)
	return e
}

func (t *Tests) SetUpTest(c *C) {
	t.LoggingSuite.SetUpTest(c)
	t.Env = t.Open(c)
}

func (t *Tests) TearDownTest(c *C) {
	if t.Env != nil {
		err := t.Env.Destroy(nil)
		c.Check(err, IsNil)
		t.Env = nil
	}
	t.LoggingSuite.TearDownTest(c)
}

func (t *Tests) TestBootstrapWithoutAdminSecret(c *C) {
	m := t.Env.Config().AllAttrs()
	delete(m, "admin-secret")
	env, err := environs.NewFromAttrs(m)
	c.Assert(err, IsNil)
	err = environs.Bootstrap(env, false, panicWrite)
	c.Assert(err, ErrorMatches, ".*admin-secret is required for bootstrap")
}

func (t *Tests) TestProviderAssignmentPolicy(c *C) {
	e := t.Open(c)
	policy := e.AssignmentPolicy()
	c.Assert(policy, FitsTypeOf, state.AssignUnused)
}

func (t *Tests) TestStartStop(c *C) {
	e := t.Open(c)

	insts, err := e.Instances(nil)
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 0)

	inst0, err := e.StartInstance("0", testing.InvalidStateInfo("0"), testing.InvalidAPIInfo("0"), nil)
	c.Assert(err, IsNil)
	c.Assert(inst0, NotNil)
	id0 := inst0.Id()

	inst1, err := e.StartInstance("1", testing.InvalidStateInfo("1"), testing.InvalidAPIInfo("1"), nil)
	c.Assert(err, IsNil)
	c.Assert(inst1, NotNil)
	id1 := inst1.Id()

	insts, err = e.Instances([]state.InstanceId{id0, id1})
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Equals, id0)
	c.Assert(insts[1].Id(), Equals, id1)

	// order of results is not specified
	insts, err = e.AllInstances()
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Not(Equals), insts[1].Id())

	err = e.StopInstances([]environs.Instance{inst0})
	c.Assert(err, IsNil)

	insts, err = e.Instances([]state.InstanceId{id0, id1})
	c.Assert(err, Equals, environs.ErrPartialInstances)
	c.Assert(insts[0], IsNil)
	c.Assert(insts[1].Id(), Equals, id1)

	insts, err = e.AllInstances()
	c.Assert(err, IsNil)
	c.Assert(insts[0].Id(), Equals, id1)
}

func (t *Tests) TestBootstrap(c *C) {
	// TODO tests for Bootstrap(true)
	e := t.Open(c)
	err := environs.Bootstrap(e, false, panicWrite)
	c.Assert(err, IsNil)

	info, apiInfo, err := e.StateInfo()
	c.Check(info.Addrs, Not(HasLen), 0)
	c.Check(apiInfo.Addrs, Not(HasLen), 0)

	err = environs.Bootstrap(e, false, panicWrite)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	e2 := t.Open(c)
	err = environs.Bootstrap(e2, false, panicWrite)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	info2, apiInfo2, err := e2.StateInfo()
	c.Check(info2, DeepEquals, info)
	c.Check(apiInfo2, DeepEquals, apiInfo)

	err = e2.Destroy(nil)
	c.Assert(err, IsNil)

	// Open again because Destroy invalidates old environments.
	e3 := t.Open(c)

	err = environs.Bootstrap(e3, false, panicWrite)
	c.Assert(err, IsNil)

	err = environs.Bootstrap(e3, false, panicWrite)
	c.Assert(err, NotNil)
}

var noRetry = trivial.AttemptStrategy{}

func (t *Tests) TestPersistence(c *C) {
	storage := t.Open(c).Storage()

	names := []string{
		"aa",
		"zzz/aa",
		"zzz/bb",
	}
	for _, name := range names {
		checkFileDoesNotExist(c, storage, name, noRetry)
		checkPutFile(c, storage, name, []byte(name))
	}
	checkList(c, storage, "", names)
	checkList(c, storage, "a", []string{"aa"})
	checkList(c, storage, "zzz/", []string{"zzz/aa", "zzz/bb"})

	storage2 := t.Open(c).Storage()
	for _, name := range names {
		checkFileHasContents(c, storage2, name, []byte(name), noRetry)
	}

	// remove the first file and check that the others remain.
	err := storage2.Remove(names[0])
	c.Check(err, IsNil)

	// check that it's ok to remove a file twice.
	err = storage2.Remove(names[0])
	c.Check(err, IsNil)

	// ... and check it's been removed in the other environment
	checkFileDoesNotExist(c, storage, names[0], noRetry)

	// ... and that the rest of the files are still around
	checkList(c, storage2, "", names[1:])

	for _, name := range names[1:] {
		err := storage2.Remove(name)
		c.Assert(err, IsNil)
	}

	// check they've all gone
	checkList(c, storage2, "", nil)
}

func checkList(c *C, storage environs.StorageReader, prefix string, names []string) {
	lnames, err := storage.List(prefix)
	c.Assert(err, IsNil)
	c.Assert(lnames, DeepEquals, names)
}

func checkPutFile(c *C, storage environs.StorageWriter, name string, contents []byte) {
	err := storage.Put(name, bytes.NewBuffer(contents), int64(len(contents)))
	c.Assert(err, IsNil)
}

func checkFileDoesNotExist(c *C, storage environs.StorageReader, name string, attempt trivial.AttemptStrategy) {
	var r io.ReadCloser
	var err error
	for a := attempt.Start(); a.Next(); {
		r, err = storage.Get(name)
		if err != nil {
			break
		}
	}
	c.Assert(r, IsNil)
	var notFoundError *environs.NotFoundError
	c.Assert(err, FitsTypeOf, notFoundError)
}

func checkFileHasContents(c *C, storage environs.StorageReader, name string, contents []byte, attempt trivial.AttemptStrategy) {
	r, err := storage.Get(name)
	c.Assert(err, IsNil)
	c.Check(r, NotNil)
	defer r.Close()

	data, err := ioutil.ReadAll(r)
	c.Check(err, IsNil)
	c.Check(data, DeepEquals, contents)

	url, err := storage.URL(name)
	c.Assert(err, IsNil)

	var resp *http.Response
	for a := attempt.Start(); a.Next(); {
		resp, err = http.Get(url)
		c.Assert(err, IsNil)
		if resp.StatusCode != 404 {
			break
		}
		c.Logf("get retrying after earlier get succeeded. *sigh*.")
	}
	c.Assert(err, IsNil)
	data, err = ioutil.ReadAll(resp.Body)
	c.Assert(err, IsNil)
	defer resp.Body.Close()
	c.Assert(resp.StatusCode, Equals, 200, Commentf("error response: %s", data))
	c.Check(data, DeepEquals, contents)
}
<commit_msg>always sort arrays inside checkList<commit_after>package jujutest

import (
	"bytes"
	"io"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/trivial"
	"net/http"
	"sort"
)

// Tests is a gocheck suite containing tests verifying juju functionality
// against the environment with the given configuration. The
// tests are not designed to be run against a live server - the Environ
// is opened once for each test, and some potentially expensive operations
// may be executed.
type Tests struct {
	coretesting.LoggingSuite
	Config map[string]interface{}
	Env    environs.Environ
}

// Open opens an instance of the testing environment.
func (t *Tests) Open(c *C) environs.Environ {
	e, err := environs.NewFromAttrs(t.Config)
	c.Assert(err, IsNil, Commentf("opening environ %#v", t.Config))
	c.Assert(e, NotNil)
	return e
}

func (t *Tests) SetUpTest(c *C) {
	t.LoggingSuite.SetUpTest(c)
	t.Env = t.Open(c)
}

func (t *Tests) TearDownTest(c *C) {
	if t.Env != nil {
		err := t.Env.Destroy(nil)
		c.Check(err, IsNil)
		t.Env = nil
	}
	t.LoggingSuite.TearDownTest(c)
}

func (t *Tests) TestBootstrapWithoutAdminSecret(c *C) {
	m := t.Env.Config().AllAttrs()
	delete(m, "admin-secret")
	env, err := environs.NewFromAttrs(m)
	c.Assert(err, IsNil)
	err = environs.Bootstrap(env, false, panicWrite)
	c.Assert(err, ErrorMatches, ".*admin-secret is required for bootstrap")
}

func (t *Tests) TestProviderAssignmentPolicy(c *C) {
	e := t.Open(c)
	policy := e.AssignmentPolicy()
	c.Assert(policy, FitsTypeOf, state.AssignUnused)
}

func (t *Tests) TestStartStop(c *C) {
	e := t.Open(c)

	insts, err := e.Instances(nil)
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 0)

	inst0, err := e.StartInstance("0", testing.InvalidStateInfo("0"), testing.InvalidAPIInfo("0"), nil)
	c.Assert(err, IsNil)
	c.Assert(inst0, NotNil)
	id0 := inst0.Id()

	inst1, err := e.StartInstance("1", testing.InvalidStateInfo("1"), testing.InvalidAPIInfo("1"), nil)
	c.Assert(err, IsNil)
	c.Assert(inst1, NotNil)
	id1 := inst1.Id()

	insts, err = e.Instances([]state.InstanceId{id0, id1})
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Equals, id0)
	c.Assert(insts[1].Id(), Equals, id1)

	// order of results is not specified
	insts, err = e.AllInstances()
	c.Assert(err, IsNil)
	c.Assert(insts, HasLen, 2)
	c.Assert(insts[0].Id(), Not(Equals), insts[1].Id())

	err = e.StopInstances([]environs.Instance{inst0})
	c.Assert(err, IsNil)

	insts, err = e.Instances([]state.InstanceId{id0, id1})
	c.Assert(err, Equals, environs.ErrPartialInstances)
	c.Assert(insts[0], IsNil)
	c.Assert(insts[1].Id(), Equals, id1)

	insts, err = e.AllInstances()
	c.Assert(err, IsNil)
	c.Assert(insts[0].Id(), Equals, id1)
}

func (t *Tests) TestBootstrap(c *C) {
	// TODO tests for Bootstrap(true)
	e := t.Open(c)
	err := environs.Bootstrap(e, false, panicWrite)
	c.Assert(err, IsNil)

	info, apiInfo, err := e.StateInfo()
	c.Check(info.Addrs, Not(HasLen), 0)
	c.Check(apiInfo.Addrs, Not(HasLen), 0)

	err = environs.Bootstrap(e, false, panicWrite)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	e2 := t.Open(c)
	err = environs.Bootstrap(e2, false, panicWrite)
	c.Assert(err, ErrorMatches, "environment is already bootstrapped")

	info2, apiInfo2, err := e2.StateInfo()
	c.Check(info2, DeepEquals, info)
	c.Check(apiInfo2, DeepEquals, apiInfo)

	err = e2.Destroy(nil)
	c.Assert(err, IsNil)

	// Open again because Destroy invalidates old environments.
	e3 := t.Open(c)

	err = environs.Bootstrap(e3, false, panicWrite)
	c.Assert(err, IsNil)

	err = environs.Bootstrap(e3, false, panicWrite)
	c.Assert(err, NotNil)
}

var noRetry = trivial.AttemptStrategy{}

func (t *Tests) TestPersistence(c *C) {
	storage := t.Open(c).Storage()

	names := []string{
		"aa",
		"zzz/aa",
		"zzz/bb",
	}
	for _, name := range names {
		checkFileDoesNotExist(c, storage, name, noRetry)
		checkPutFile(c, storage, name, []byte(name))
	}
	checkList(c, storage, "", names)
	checkList(c, storage, "a", []string{"aa"})
	checkList(c, storage, "zzz/", []string{"zzz/aa", "zzz/bb"})

	storage2 := t.Open(c).Storage()
	for _, name := range names {
		checkFileHasContents(c, storage2, name, []byte(name), noRetry)
	}

	// remove the first file and check that the others remain.
	err := storage2.Remove(names[0])
	c.Check(err, IsNil)

	// check that it's ok to remove a file twice.
	err = storage2.Remove(names[0])
	c.Check(err, IsNil)

	// ... and check it's been removed in the other environment
	checkFileDoesNotExist(c, storage, names[0], noRetry)

	// ... and that the rest of the files are still around
	checkList(c, storage2, "", names[1:])

	for _, name := range names[1:] {
		err := storage2.Remove(name)
		c.Assert(err, IsNil)
	}

	// check they've all gone
	checkList(c, storage2, "", nil)
}

func checkList(c *C, storage environs.StorageReader, prefix string, names []string) {
	lnames, err := storage.List(prefix)
	c.Assert(err, IsNil)
	// TODO(dfc) gocheck should grow an ArrayEquals checker.
	expected := copyslice(lnames)
	sort.Strings(expected)
	actual := copyslice(names)
	sort.Strings(actual)
	c.Assert(expected, DeepEquals, actual)
}

// copyslice returns a copy of the slice
func copyslice(s []string) []string {
	r := make([]string, len(s))
	copy(r, s)
	return r
}

func checkPutFile(c *C, storage environs.StorageWriter, name string, contents []byte) {
	err := storage.Put(name, bytes.NewBuffer(contents), int64(len(contents)))
	c.Assert(err, IsNil)
}

func checkFileDoesNotExist(c *C, storage environs.StorageReader, name string, attempt trivial.AttemptStrategy) {
	var r io.ReadCloser
	var err error
	for a := attempt.Start(); a.Next(); {
		r, err = storage.Get(name)
		if err != nil {
			break
		}
	}
	c.Assert(r, IsNil)
	var notFoundError *environs.NotFoundError
	c.Assert(err, FitsTypeOf, notFoundError)
}

func checkFileHasContents(c *C, storage environs.StorageReader, name string, contents []byte, attempt trivial.AttemptStrategy) {
	r, err := storage.Get(name)
	c.Assert(err, IsNil)
	c.Check(r, NotNil)
	defer r.Close()

	data, err := ioutil.ReadAll(r)
	c.Check(err, IsNil)
	c.Check(data, DeepEquals, contents)

	url, err := storage.URL(name)
	c.Assert(err, IsNil)

	var resp *http.Response
	for a := attempt.Start(); a.Next(); {
		resp, err = http.Get(url)
		c.Assert(err, IsNil)
		if resp.StatusCode != 404 {
			break
		}
		c.Logf("get retrying after earlier get succeeded. *sigh*.")
	}
	c.Assert(err, IsNil)
	data, err = ioutil.ReadAll(resp.Body)
	c.Assert(err, IsNil)
	defer resp.Body.Close()
	c.Assert(resp.StatusCode, Equals, 200, Commentf("error response: %s", data))
	c.Check(data, DeepEquals, contents)
}
<|endoftext|>


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

package scripting

import (
	"testing"

	"github.com/OWASP/Amass/v3/requests"
	"github.com/caffix/stringset"
)

func TestNewNames(t *testing.T) {
	expected := []string{"owasp.org", "www.owasp.org", "ftp.owasp.org", "mail.owasp.org",
		"dev.owasp.org", "prod.owasp.org", "vpn.owasp.org", "uat.owasp.org", "stage.owasp.org",
		"confluence.owasp.org", "api.owasp.org", "test.owasp.org"}

	ctx, sys := setupMockScriptEnv(`
		name="names"
		type="testing"

		function vertical(ctx, domain)
			new_name(ctx, "owasp.org")

			local content = [[
				www.owasp.org
				www.owasp.org
				www.owasp.org
				www.owasp.org
				www.owasp.org
				ftp.owasp.org
				mail.owasp.org
				dev.owasp.org
				prod.owasp.org
				vpn.owasp.org
				uat.owasp.org
				stage.owasp.org
				confluence.owasp.org
				api.owasp.org
				test.owasp.org
			]]
			send_names(ctx, content)
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.DNSRequest, num)
	fn := func(req *requests.DNSRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewNameTopic, fn)
	defer bus.Unsubscribe(requests.NewNameTopic, fn)

	domain := "owasp.org"
	cfg.AddDomain(domain)
	sys.DataSources()[0].Request(ctx, &requests.DNSRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp := expected[i]; req.Name != exp ||
			req.Domain != domain || req.Tag != "testing" || req.Source != "names" {
			t.Errorf("Incorrect output for name %d, expected: %s, got: %v", i+1, exp, req)
		}
	}
}

func TestNewAddrs(t *testing.T) {
	expected := []string{"72.237.4.113", "72.237.4.114", "72.237.4.35", "72.237.4.38", "72.237.4.79",
		"72.237.4.90", "72.237.4.103", "72.237.4.243", "4.26.24.234", "44.193.34.238", "52.206.190.41", "18.211.32.87"}

	ctx, sys := setupMockScriptEnv(`
		name="addrs"
		type="testing"

		function vertical(ctx, domain)
			local addrs = {"72.237.4.113", "72.237.4.114", "72.237.4.35",
				"72.237.4.38", "72.237.4.79", "72.237.4.90", "72.237.4.103",
				"72.237.4.243", "4.26.24.234", "44.193.34.238", "52.206.190.41", "18.211.32.87"}

			for _, addr in ipairs(addrs) do
				new_addr(ctx, addr, domain)
    		end
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.AddrRequest, num)
	fn := func(req *requests.AddrRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewAddrTopic, fn)
	defer bus.Unsubscribe(requests.NewAddrTopic, fn)

	domain := "owasp.org"
	cfg.AddDomain(domain)
	sys.DataSources()[0].Request(ctx, &requests.DNSRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp := expected[i]; req.Address != exp ||
			req.Domain != domain || req.Tag != "testing" || req.Source != "addrs" {
			t.Errorf("Incorrect output for address %d, expected: %s, got: %v", i+1, exp, req)
		}
	}
}

func TestNewASNs(t *testing.T) {
	expected := []requests.ASNRequest{
		{
			Address:     "72.237.4.113",
			ASN:         26808,
			Prefix:      "72.237.4.0/24",
			Description: "UTICA-COLLEGE - Utica College",
			Netblocks:   stringset.New("72.237.4.0/24"),
			Tag:         "testing",
			Source:      "asns",
		},
		{
			Address:     "104.16.0.1",
			ASN:         13335,
			Prefix:      "104.16.0.0/14",
			CC:          "US",
			Registry:    "ARIN",
			Description: "CLOUDFLARENET - Cloudflare, Inc.",
			Netblocks:   stringset.New("104.16.0.0/14", "2606:4700::/47"),
			Tag:         "testing",
			Source:      "asns",
		},
	}

	ctx, sys := setupMockScriptEnv(`
		name="asns"
		type="testing"

		function asn(ctx, addr, asn)
			new_asn(ctx, {
				addr="72.237.4.113",
				['asn']=26808,
				prefix="72.237.4.0/24",
				desc="UTICA-COLLEGE - Utica College",
				netblocks={"72.237.4.0/24"},
			})

			new_asn(ctx, {
				addr="104.16.0.1",
				['asn']=13335,
				prefix="104.16.0.0/14",
				cc="US",
				registry="ARIN",
				desc="CLOUDFLARENET - Cloudflare, Inc.",
				netblocks={
					"104.16.0.0/14",
					"2606:4700::/47",
				},
			})

			new_asn(ctx, {
				addr="not.a.valid.addr",
				['asn']=15169,
				prefix="172.217.0.0/19",
				desc="GOOGLE - Google LLC",
				netblocks={
					"172.217.0.0/19",
					"2607:f8b0:4004::/48",
				},
			})

			new_asn(ctx, {
				addr="52.8.0.1",
				['asn']=16509,
				prefix="52.8.0.0/invalid",
				desc="AMAZON-02 - Amazon.com, Inc.",
				netblocks={
					"52.8.0.0/13",
					"50.18.0.0/16",
				},
			})

			new_asn(ctx, {
				addr="162.242.128.1",
				['asn']=33070,
				prefix="162.242.128.0/19",
				desc="",
				netblocks={"162.242.128.0/19"},
			})
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the event bus")
	}

	num := len(expected)
	ch := make(chan *requests.ASNRequest, num)
	fn := func(req *requests.ASNRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewASNTopic, fn)
	defer bus.Unsubscribe(requests.NewASNTopic, fn)

	address := "72.237.4.113"
	sys.DataSources()[0].Request(ctx, &requests.ASNRequest{Address: address})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp := expected[i]; !matchingASNs(req, &exp) {
			t.Errorf("Incorrect output for ASN %d, expected: %v, got: %v", i+1, exp, req)
		}
	}
}

func matchingASNs(first, second *requests.ASNRequest) bool {
	match := true

	if addr := first.Address; addr == "" || addr != second.Address {
		match = false
	}
	if asn := first.ASN; asn == 0 || asn != second.ASN {
		match = false
	}
	if prefix := first.Prefix; prefix == "" || prefix != second.Prefix {
		match = false
	}
	if cc := first.CC; cc != "" && cc != second.CC {
		match = false
	}
	if reg := first.Registry; reg != "" && reg != second.Registry {
		match = false
	}
	if desc := first.Description; desc == "" || desc != second.Description {
		match = false
	}

	nb := first.Netblocks
	if nb.Len() == 0 {
		match = false
	}

	nb.Subtract(second.Netblocks)
	if nb.Len() != 0 {
		match = false
	}

	if tag := first.Tag; tag == "" || tag != second.Tag {
		match = false
	}
	if src := first.Source; src == "" || src != second.Source {
		match = false
	}

	return match
}

func TestAssociated(t *testing.T) {
	expected := []requests.WhoisRequest{
		{
			Domain:     "owasp.org",
			NewDomains: []string{"globalappsec.org"},
			Tag:        "testing",
			Source:     "associated",
		},
		{
			Domain:     "utica.edu",
			NewDomains: []string{"necyber.com"},
			Tag:        "testing",
			Source:     "associated",
		},
	}

	ctx, sys := setupMockScriptEnv(`
		name="associated"
		type="testing"

		function horizontal(ctx, domain)
			local assocs = {
				{
					Domain="owasp.org",
					Assoc="globalappsec.org",
				},
				{
					Domain="utica.edu",
					Assoc="necyber.com",
				},
			}

			for _, a in ipairs(assocs) do
				associated(ctx, a.Domain, a.Assoc)
    		end
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.WhoisRequest, num)
	fn := func(req *requests.WhoisRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewWhoisTopic, fn)
	defer bus.Unsubscribe(requests.NewWhoisTopic, fn)

	domain := "owasp.org"
	cfg.AddDomains(domain, "utica.edu")
	sys.DataSources()[0].Request(ctx, &requests.WhoisRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp := expected[i]; req.Domain != exp.Domain ||
			req.NewDomains[0] != exp.NewDomains[0] || req.Tag != "testing" || req.Source != "associated" {
			t.Errorf("Incorrect output for associated %d, expected: %s, got: %v", i+1, exp, req)
		}
	}
}
<commit_msg>Improved the unit tests for new data functions<commit_after>// Copyright 2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package scripting

import (
	"testing"

	"github.com/OWASP/Amass/v3/requests"
	"github.com/caffix/stringset"
)

func TestNewNames(t *testing.T) {
	expected := stringset.New("owasp.org", "www.owasp.org", "ftp.owasp.org", "mail.owasp.org",
		"dev.owasp.org", "prod.owasp.org", "vpn.owasp.org", "uat.owasp.org", "stage.owasp.org",
		"confluence.owasp.org", "api.owasp.org", "test.owasp.org")

	ctx, sys := setupMockScriptEnv(`
		name="names"
		type="testing"

		function vertical(ctx, domain)
			new_name(ctx, "owasp.org")

			local content = [[
				www.owasp.org
				www.owasp.org
				www.owasp.org
				www.owasp.org
				www.owasp.org
				ftp.owasp.org
				mail.owasp.org
				dev.owasp.org
				prod.owasp.org
				vpn.owasp.org
				uat.owasp.org
				stage.owasp.org
				confluence.owasp.org
				api.owasp.org
				test.owasp.org
			]]
			send_names(ctx, content)
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.DNSRequest, num)
	fn := func(req *requests.DNSRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewNameTopic, fn)
	defer bus.Unsubscribe(requests.NewNameTopic, fn)

	domain := "owasp.org"
	cfg.AddDomain(domain)
	sys.DataSources()[0].Request(ctx, &requests.DNSRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if !expected.Has(req.Name) || req.Domain != domain || req.Tag != "testing" || req.Source != "names" {
			t.Errorf("Name %d: %v was not found in the list of expected names", i+1, req.Name)
		}

		expected.Remove(req.Name)
	}
}

func TestNewAddrs(t *testing.T) {
	expected := stringset.New("72.237.4.113", "72.237.4.114", "72.237.4.35", "72.237.4.38", "72.237.4.79",
		"72.237.4.90", "72.237.4.103", "72.237.4.243", "4.26.24.234", "44.193.34.238", "52.206.190.41", "18.211.32.87")

	ctx, sys := setupMockScriptEnv(`
		name="addrs"
		type="testing"

		function vertical(ctx, domain)
			local addrs = {"72.237.4.113", "72.237.4.114", "72.237.4.35",
				"72.237.4.38", "72.237.4.79", "72.237.4.90", "72.237.4.103",
				"72.237.4.243", "4.26.24.234", "44.193.34.238", "52.206.190.41", "18.211.32.87"}

			for _, addr in ipairs(addrs) do
				new_addr(ctx, addr, domain)
    		end
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.AddrRequest, num)
	fn := func(req *requests.AddrRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewAddrTopic, fn)
	defer bus.Unsubscribe(requests.NewAddrTopic, fn)

	domain := "owasp.org"
	cfg.AddDomain(domain)
	sys.DataSources()[0].Request(ctx, &requests.DNSRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if !expected.Has(req.Address) || req.Domain != domain || req.Tag != "testing" || req.Source != "addrs" {
			t.Errorf("Address %d: %v was not found in the list of expected addresses", i+1, req.Address)
		}

		expected.Remove(req.Address)
	}
}

func TestNewASNs(t *testing.T) {
	expected := map[int]*requests.ASNRequest{
		26808: {
			Address:     "72.237.4.113",
			ASN:         26808,
			Prefix:      "72.237.4.0/24",
			Description: "UTICA-COLLEGE - Utica College",
			Netblocks:   stringset.New("72.237.4.0/24"),
			Tag:         "testing",
			Source:      "asns",
		},
		13335: {
			Address:     "104.16.0.1",
			ASN:         13335,
			Prefix:      "104.16.0.0/14",
			CC:          "US",
			Registry:    "ARIN",
			Description: "CLOUDFLARENET - Cloudflare, Inc.",
			Netblocks:   stringset.New("104.16.0.0/14", "2606:4700::/47"),
			Tag:         "testing",
			Source:      "asns",
		},
	}

	ctx, sys := setupMockScriptEnv(`
		name="asns"
		type="testing"

		function asn(ctx, addr, asn)
			new_asn(ctx, {
				addr="72.237.4.113",
				['asn']=26808,
				prefix="72.237.4.0/24",
				desc="UTICA-COLLEGE - Utica College",
				netblocks={"72.237.4.0/24"},
			})

			new_asn(ctx, {
				addr="104.16.0.1",
				['asn']=13335,
				prefix="104.16.0.0/14",
				cc="US",
				registry="ARIN",
				desc="CLOUDFLARENET - Cloudflare, Inc.",
				netblocks={
					"104.16.0.0/14",
					"2606:4700::/47",
				},
			})

			new_asn(ctx, {
				addr="not.a.valid.addr",
				['asn']=15169,
				prefix="172.217.0.0/19",
				desc="GOOGLE - Google LLC",
				netblocks={
					"172.217.0.0/19",
					"2607:f8b0:4004::/48",
				},
			})

			new_asn(ctx, {
				addr="52.8.0.1",
				['asn']=16509,
				prefix="52.8.0.0/invalid",
				desc="AMAZON-02 - Amazon.com, Inc.",
				netblocks={
					"52.8.0.0/13",
					"50.18.0.0/16",
				},
			})

			new_asn(ctx, {
				addr="162.242.128.1",
				['asn']=33070,
				prefix="162.242.128.0/19",
				desc="",
				netblocks={"162.242.128.0/19"},
			})
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	_, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the event bus")
	}

	num := len(expected)
	ch := make(chan *requests.ASNRequest, num)
	fn := func(req *requests.ASNRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewASNTopic, fn)
	defer bus.Unsubscribe(requests.NewASNTopic, fn)

	address := "72.237.4.113"
	sys.DataSources()[0].Request(ctx, &requests.ASNRequest{Address: address})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp, found := expected[req.ASN]; !found || !matchingASNs(req, exp) {
			t.Errorf("Incorrect output for ASN %d, expected: %v, got: %v", i+1, exp, req)
		}
	}
}

func matchingASNs(first, second *requests.ASNRequest) bool {
	match := true

	if addr := first.Address; addr == "" || addr != second.Address {
		match = false
	}
	if asn := first.ASN; asn == 0 || asn != second.ASN {
		match = false
	}
	if prefix := first.Prefix; prefix == "" || prefix != second.Prefix {
		match = false
	}
	if cc := first.CC; cc != "" && cc != second.CC {
		match = false
	}
	if reg := first.Registry; reg != "" && reg != second.Registry {
		match = false
	}
	if desc := first.Description; desc == "" || desc != second.Description {
		match = false
	}

	nb := first.Netblocks
	if nb.Len() == 0 {
		match = false
	}

	nb.Subtract(second.Netblocks)
	if nb.Len() != 0 {
		match = false
	}

	if tag := first.Tag; tag == "" || tag != second.Tag {
		match = false
	}
	if src := first.Source; src == "" || src != second.Source {
		match = false
	}

	return match
}

func TestAssociated(t *testing.T) {
	expected := map[string]*requests.WhoisRequest{
		"owasp.org": {
			Domain:     "owasp.org",
			NewDomains: []string{"globalappsec.org"},
			Tag:        "testing",
			Source:     "associated",
		},
		"utica.edu": {
			Domain:     "utica.edu",
			NewDomains: []string{"necyber.com"},
			Tag:        "testing",
			Source:     "associated",
		},
	}

	ctx, sys := setupMockScriptEnv(`
		name="associated"
		type="testing"

		function horizontal(ctx, domain)
			local assocs = {
				{
					Domain="owasp.org",
					Assoc="globalappsec.org",
				},
				{
					Domain="utica.edu",
					Assoc="necyber.com",
				},
			}

			for _, a in ipairs(assocs) do
				associated(ctx, a.Domain, a.Assoc)
    		end
		end
	`)
	if ctx == nil || sys == nil {
		t.Fatal("Failed to initialize the scripting environment")
	}
	defer func() { _ = sys.Shutdown() }()

	cfg, bus, err := requests.ContextConfigBus(ctx)
	if err != nil {
		t.Fatal("Failed to obtain the config and event bus")
	}

	num := len(expected)
	ch := make(chan *requests.WhoisRequest, num)
	fn := func(req *requests.WhoisRequest) {
		ch <- req
	}

	bus.Subscribe(requests.NewWhoisTopic, fn)
	defer bus.Unsubscribe(requests.NewWhoisTopic, fn)

	domain := "owasp.org"
	cfg.AddDomains(domain, "utica.edu")
	sys.DataSources()[0].Request(ctx, &requests.WhoisRequest{Domain: domain})

	for i := 0; i < num; i++ {
		req := <-ch

		if exp, found := expected[req.Domain]; !found || req.Domain != exp.Domain ||
			req.NewDomains[0] != exp.NewDomains[0] || req.Tag != "testing" || req.Source != "associated" {
			t.Errorf("Incorrect output for associated %d, expected: %s, got: %v", i+1, exp, req)
		}
	}
}
<|endoftext|>


<commit_before>package cloudformation

type UpdateStackParameters struct {
	BaseParameters
	StackPolicyDuringUpdateBody string
	StackPolicyDuringUpdateURL  string
}

func (p *UpdateStackParameters) values() Values {
	v := p.BaseParameters.values()
	v["StackPolicyDuringUpdateBody"] = p.StackPolicyDuringUpdateBody
	v["StackPolicyDuringUpdateURL"] = p.StackPolicyDuringUpdateURL
	return v
}

type UpdateStackResponse struct {
	UpdateStackResult *UpdateStackResult `xml:"UpdateStackResult"`
}

type UpdateStackResult struct {
	StackId string `xml:"StackId"`
}

type UpdateStack struct {
	Parameters                  []*StackParameter
	Capabilities                []string
	StackName                   string
	StackPolicyBody             string
	StackPolicyURL              string
	TemplateBody                string
	TemplateURL                 string
	StackPolicyDuringUpdateBody string
	StackPolicyDuringUpdateURL  string
}

func (update *UpdateStack) values() Values {
	v := Values{
		"StackName":                   update.StackName,
		"StackPolicyBody":             update.StackPolicyBody,
		"StackPolicyURL":              update.StackPolicyURL,
		"TemplateBody":                update.TemplateBody,
		"TemplateURL":                 update.TemplateURL,
		"StackPolicyDuringUpdateBody": update.StackPolicyDuringUpdateBody,
		"StackPolicyDuringUpdateURL":  update.StackPolicyDuringUpdateURL,
	}
	v.updateCapabilities(update.Capabilities)
	v.updateParameters(update.Parameters)
	return v
}

func (update *UpdateStack) Execute(client *Client) (*UpdateStackResponse, error) {
	r := &UpdateStackResponse{}
	e := client.loadCloudFormationResource("UpdateStack", update.values(), r)
	return r, e
}

func (c *Client) UpdateStack(params UpdateStackParameters) (stackId string, e error) {
	r := &UpdateStackResponse{}
	e = c.loadCloudFormationResource("UpdateStack", params.values(), r)
	if e != nil {
		return "", e
	}
	return r.UpdateStackResult.StackId, nil
}
<commit_msg>use dedicated error for cloudformation "no updates" scenarion<commit_after>package cloudformation

import "fmt"

type UpdateStackParameters struct {
	BaseParameters
	StackPolicyDuringUpdateBody string
	StackPolicyDuringUpdateURL  string
}

func (p *UpdateStackParameters) values() Values {
	v := p.BaseParameters.values()
	v["StackPolicyDuringUpdateBody"] = p.StackPolicyDuringUpdateBody
	v["StackPolicyDuringUpdateURL"] = p.StackPolicyDuringUpdateURL
	return v
}

type UpdateStackResponse struct {
	UpdateStackResult *UpdateStackResult `xml:"UpdateStackResult"`
}

type UpdateStackResult struct {
	StackId string `xml:"StackId"`
}

type UpdateStack struct {
	Parameters                  []*StackParameter
	Capabilities                []string
	StackName                   string
	StackPolicyBody             string
	StackPolicyURL              string
	TemplateBody                string
	TemplateURL                 string
	StackPolicyDuringUpdateBody string
	StackPolicyDuringUpdateURL  string
}

func (update *UpdateStack) values() Values {
	v := Values{
		"StackName":                   update.StackName,
		"StackPolicyBody":             update.StackPolicyBody,
		"StackPolicyURL":              update.StackPolicyURL,
		"TemplateBody":                update.TemplateBody,
		"TemplateURL":                 update.TemplateURL,
		"StackPolicyDuringUpdateBody": update.StackPolicyDuringUpdateBody,
		"StackPolicyDuringUpdateURL":  update.StackPolicyDuringUpdateURL,
	}
	v.updateCapabilities(update.Capabilities)
	v.updateParameters(update.Parameters)
	return v
}

func (update *UpdateStack) Execute(client *Client) (*UpdateStackResponse, error) {
	r := &UpdateStackResponse{}
	e := client.loadCloudFormationResource("UpdateStack", update.values(), r)
	return r, e
}

const errorNoUpdate = "No updates are to be performed."

var ErrorNoUpdate = fmt.Errorf(errorNoUpdate)

func (c *Client) UpdateStack(params UpdateStackParameters) (stackId string, e error) {
	r := &UpdateStackResponse{}
	e = c.loadCloudFormationResource("UpdateStack", params.values(), r)
	if e != nil {
		if e.Error() == errorNoUpdate {
			return "", ErrorNoUpdate
		}
		return "", e
	}
	return r.UpdateStackResult.StackId, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"errors"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/hashicorp/terraform/helper/schema"
)

func dataSourceAwsLbListener() *schema.Resource {
	return &schema.Resource{
		Read: dataSourceAwsLbListenerRead,

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"load_balancer_arn", "port"},
			},

			"load_balancer_arn": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"arn"},
			},
			"port": {
				Type:          schema.TypeInt,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"arn"},
			},

			"protocol": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"ssl_policy": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"certificate_arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"default_action": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"target_group_arn": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"type": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
		},
	}
}

func dataSourceAwsLbListenerRead(d *schema.ResourceData, meta interface{}) error {
	if _, ok := d.GetOk("arn"); ok {
		d.SetId(d.Get("arn").(string))
		//log.Printf("[DEBUG] read listener %s", d.Get("arn").(string))
		return resourceAwsLbListenerRead(d, meta)
	}

	conn := meta.(*AWSClient).elbv2conn
	lbArn, lbOk := d.GetOk("load_balancer_arn")
	port, portOk := d.GetOk("port")
	if !lbOk || !portOk {
		return errors.New("both load_balancer_arn and port must be set")
	}
	resp, err := conn.DescribeListeners(&elbv2.DescribeListenersInput{
		LoadBalancerArn: aws.String(lbArn.(string)),
	})
	if err != nil {
		return err
	}
	if len(resp.Listeners) == 0 {
		return fmt.Errorf("[DEBUG] no listener exists for load balancer: %s", lbArn)
	}
	for _, listener := range resp.Listeners {
		if *listener.Port == int64(port.(int)) {
			//log.Printf("[DEBUG] get listener arn for %s:%s: %s", lbArn, port, *listener.Port)
			d.SetId(*listener.ListenerArn)
			return resourceAwsLbListenerRead(d, meta)
		}
	}

	return errors.New("failed to get listener arn with given arguments")
}
<commit_msg>Implement data source: aws_lb_listener<commit_after>package aws

import (
	"errors"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/hashicorp/terraform/helper/schema"
)

func dataSourceAwsLbListener() *schema.Resource {
	return &schema.Resource{
		Read: dataSourceAwsLbListenerRead,

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"load_balancer_arn", "port"},
			},

			"load_balancer_arn": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"arn"},
			},
			"port": {
				Type:          schema.TypeInt,
				Optional:      true,
				Computed:      true,
				ConflictsWith: []string{"arn"},
			},

			"protocol": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"ssl_policy": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"certificate_arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"default_action": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"type": {
							Type:     schema.TypeString,
							Computed: true,
						},

						"target_group_arn": {
							Type:     schema.TypeString,
							Computed: true,
						},

						"redirect": {
							Type:     schema.TypeList,
							Computed: true,
							MaxItems: 1,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"host": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"path": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"port": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"protocol": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"query": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"status_code": {
										Type:     schema.TypeString,
										Computed: true,
									},
								},
							},
						},

						"fixed_response": {
							Type:     schema.TypeList,
							Computed: true,
							MaxItems: 1,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"content_type": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"message_body": {
										Type:     schema.TypeString,
										Computed: true,
									},

									"status_code": {
										Type:     schema.TypeString,
										Computed: true,
									},
								},
							},
						},
					},
				},
			},
		},
	}
}

func dataSourceAwsLbListenerRead(d *schema.ResourceData, meta interface{}) error {
	if _, ok := d.GetOk("arn"); ok {
		d.SetId(d.Get("arn").(string))
		//log.Printf("[DEBUG] read listener %s", d.Get("arn").(string))
		return resourceAwsLbListenerRead(d, meta)
	}

	conn := meta.(*AWSClient).elbv2conn
	lbArn, lbOk := d.GetOk("load_balancer_arn")
	port, portOk := d.GetOk("port")
	if !lbOk || !portOk {
		return errors.New("both load_balancer_arn and port must be set")
	}
	resp, err := conn.DescribeListeners(&elbv2.DescribeListenersInput{
		LoadBalancerArn: aws.String(lbArn.(string)),
	})
	if err != nil {
		return err
	}
	if len(resp.Listeners) == 0 {
		return fmt.Errorf("[DEBUG] no listener exists for load balancer: %s", lbArn)
	}
	for _, listener := range resp.Listeners {
		if *listener.Port == int64(port.(int)) {
			//log.Printf("[DEBUG] get listener arn for %s:%s: %s", lbArn, port, *listener.Port)
			d.SetId(*listener.ListenerArn)
			return resourceAwsLbListenerRead(d, meta)
		}
	}

	return errors.New("failed to get listener arn with given arguments")
}
<|endoftext|>