<commit_before>package cmd

import (
	"fmt"
	"io/ioutil"
	"runtime"
	"strconv"
	"strings"

	"github.com/MakeNowJust/heredoc/v2"
	"github.com/rsteube/carapace"
	"github.com/spf13/cobra"
	gitlab "github.com/xanzy/go-gitlab"
	"github.com/zaquestion/lab/internal/action"
	"github.com/zaquestion/lab/internal/git"
	lab "github.com/zaquestion/lab/internal/gitlab"
)

var mrCreateDiscussionCmd = &cobra.Command{
	Use:     "discussion [remote] [<MR ID or branch>]",
	Short:   "Start a discussion on an MR on GitLab",
	Aliases: []string{"block", "thread"},
	Example: heredoc.Doc(`
		lab mr discussion origin
		lab mr discussion my_remote -m "discussion comment"
		lab mr discussion upstream -F test_file.txt
		lab mr discussion --commit abcdef123456`),
	PersistentPreRun: labPersistentPreRun,
	Run: func(cmd *cobra.Command, args []string) {
		rn, mrNum, err := parseArgsWithGitBranchMR(args)
		if err != nil {
			log.Fatal(err)
		}

		msgs, err := cmd.Flags().GetStringSlice("message")
		if err != nil {
			log.Fatal(err)
		}

		filename, err := cmd.Flags().GetString("file")
		if err != nil {
			log.Fatal(err)
		}

		commit, err := cmd.Flags().GetString("commit")
		if err != nil {
			log.Fatal(err)
		}

		state := noteGetState(rn, true, int(mrNum))

		body := ""
		if filename != "" {
			content, err := ioutil.ReadFile(filename)
			if err != nil {
				log.Fatal(err)
			}
			body = string(content)
		} else if commit == "" {
			body, err = mrDiscussionMsg(int(mrNum), state, commit, msgs, "\n")
			if err != nil {
				_, f, l, _ := runtime.Caller(0)
				log.Fatal(f+":"+strconv.Itoa(l)+" ", err)
			}
		} else {
			body = getCommitBody(rn, commit)
			body, err = mrDiscussionMsg(int(mrNum), state, commit, nil, body)
			if err != nil {
				_, f, l, _ := runtime.Caller(0)
				log.Fatal(f+":"+strconv.Itoa(l)+" ", err)
			}
			createCommitComments(rn, int(mrNum), commit, body, true)
			return
		}

		if body == "" {
			log.Fatal("aborting discussion due to empty discussion msg")
		}

		discussionURL, err := lab.MRCreateDiscussion(rn, int(mrNum), &gitlab.CreateMergeRequestDiscussionOptions{
			Body: &body,
		})
		if err != nil {
			log.Fatal(err)
		}
		fmt.Println(discussionURL)
	},
}

func mrDiscussionMsg(mrNum int, state string, commit string, msgs []string, body string) (string, error) {
	if len(msgs) > 0 {
		return strings.Join(msgs[0:], "\n\n"), nil
	}

	tmpl := mrDiscussionGetTemplate(commit)
	text, err := noteText(mrNum, state, commit, body, tmpl)
	if err != nil {
		return "", err
	}
	return git.EditFile("MR_DISCUSSION", text)
}

func mrDiscussionGetTemplate(commit string) string {
	if commit == "" {
		return heredoc.Doc(`
		{{.InitMsg}}
		{{.CommentChar}} This thread is being started on {{.State}} Merge Request {{.IDnum}}.
		{{.CommentChar}} Comment lines beginning with '{{.CommentChar}}' are discarded.`)
	}
	return heredoc.Doc(`
		{{.InitMsg}}
		{{.CommentChar}} This thread is being started on {{.State}} Merge Request {{.IDnum}} commit {{.Commit}}.
		{{.CommentChar}} Do not delete patch tracking lines that begin with '|'.
		{{.CommentChar}} Comment lines beginning with '{{.CommentChar}}' are discarded.`)
}

func init() {
	mrCreateDiscussionCmd.Flags().StringSliceP("message", "m", []string{}, "use the given <msg>; multiple -m are concatenated as separate paragraphs")
	mrCreateDiscussionCmd.Flags().StringP("file", "F", "", "use the given file as the message")
	mrCreateDiscussionCmd.Flags().StringP("commit", "c", "", "start a thread on a commit")

	mrCmd.AddCommand(mrCreateDiscussionCmd)
	carapace.Gen(mrCreateDiscussionCmd).PositionalCompletion(
		action.Remotes(),
		action.MergeRequests(mrList),
	)
}
<commit_msg>mr discussion: add some examples for optional MR ID / branch<commit_after>package cmd

import (
	"fmt"
	"io/ioutil"
	"runtime"
	"strconv"
	"strings"

	"github.com/MakeNowJust/heredoc/v2"
	"github.com/rsteube/carapace"
	"github.com/spf13/cobra"
	gitlab "github.com/xanzy/go-gitlab"
	"github.com/zaquestion/lab/internal/action"
	"github.com/zaquestion/lab/internal/git"
	lab "github.com/zaquestion/lab/internal/gitlab"
)

var mrCreateDiscussionCmd = &cobra.Command{
	Use:     "discussion [remote] [<MR ID or branch>]",
	Short:   "Start a discussion on an MR on GitLab",
	Aliases: []string{"block", "thread"},
	Example: heredoc.Doc(`
		lab mr discussion
		lab mr discussion origin
		lab mr discussion my_remote -m "discussion comment"
		lab mr discussion upstream -F test_file.txt
		lab mr discussion --commit abcdef123456
		lab mr discussion my-topic-branch
		lab mr discussion origin 123
		lab mr discussion origin my-topic-branch
		`),
	PersistentPreRun: labPersistentPreRun,
	Run: func(cmd *cobra.Command, args []string) {
		rn, mrNum, err := parseArgsWithGitBranchMR(args)
		if err != nil {
			log.Fatal(err)
		}

		msgs, err := cmd.Flags().GetStringSlice("message")
		if err != nil {
			log.Fatal(err)
		}

		filename, err := cmd.Flags().GetString("file")
		if err != nil {
			log.Fatal(err)
		}

		commit, err := cmd.Flags().GetString("commit")
		if err != nil {
			log.Fatal(err)
		}

		state := noteGetState(rn, true, int(mrNum))

		body := ""
		if filename != "" {
			content, err := ioutil.ReadFile(filename)
			if err != nil {
				log.Fatal(err)
			}
			body = string(content)
		} else if commit == "" {
			body, err = mrDiscussionMsg(int(mrNum), state, commit, msgs, "\n")
			if err != nil {
				_, f, l, _ := runtime.Caller(0)
				log.Fatal(f+":"+strconv.Itoa(l)+" ", err)
			}
		} else {
			body = getCommitBody(rn, commit)
			body, err = mrDiscussionMsg(int(mrNum), state, commit, nil, body)
			if err != nil {
				_, f, l, _ := runtime.Caller(0)
				log.Fatal(f+":"+strconv.Itoa(l)+" ", err)
			}
			createCommitComments(rn, int(mrNum), commit, body, true)
			return
		}

		if body == "" {
			log.Fatal("aborting discussion due to empty discussion msg")
		}

		discussionURL, err := lab.MRCreateDiscussion(rn, int(mrNum), &gitlab.CreateMergeRequestDiscussionOptions{
			Body: &body,
		})
		if err != nil {
			log.Fatal(err)
		}
		fmt.Println(discussionURL)
	},
}

func mrDiscussionMsg(mrNum int, state string, commit string, msgs []string, body string) (string, error) {
	if len(msgs) > 0 {
		return strings.Join(msgs[0:], "\n\n"), nil
	}

	tmpl := mrDiscussionGetTemplate(commit)
	text, err := noteText(mrNum, state, commit, body, tmpl)
	if err != nil {
		return "", err
	}
	return git.EditFile("MR_DISCUSSION", text)
}

func mrDiscussionGetTemplate(commit string) string {
	if commit == "" {
		return heredoc.Doc(`
		{{.InitMsg}}
		{{.CommentChar}} This thread is being started on {{.State}} Merge Request {{.IDnum}}.
		{{.CommentChar}} Comment lines beginning with '{{.CommentChar}}' are discarded.`)
	}
	return heredoc.Doc(`
		{{.InitMsg}}
		{{.CommentChar}} This thread is being started on {{.State}} Merge Request {{.IDnum}} commit {{.Commit}}.
		{{.CommentChar}} Do not delete patch tracking lines that begin with '|'.
		{{.CommentChar}} Comment lines beginning with '{{.CommentChar}}' are discarded.`)
}

func init() {
	mrCreateDiscussionCmd.Flags().StringSliceP("message", "m", []string{}, "use the given <msg>; multiple -m are concatenated as separate paragraphs")
	mrCreateDiscussionCmd.Flags().StringP("file", "F", "", "use the given file as the message")
	mrCreateDiscussionCmd.Flags().StringP("commit", "c", "", "start a thread on a commit")

	mrCmd.AddCommand(mrCreateDiscussionCmd)
	carapace.Gen(mrCreateDiscussionCmd).PositionalCompletion(
		action.Remotes(),
		action.MergeRequests(mrList),
	)
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"net"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/juju/ratelimit"
	"github.com/syncthing/syncthing/lib/relay/protocol"
	"github.com/syncthing/syncthing/lib/tlsutil"

	syncthingprotocol "github.com/syncthing/syncthing/lib/protocol"
)

var (
	listen string
	debug  bool = false

	sessionAddress []byte
	sessionPort    uint16

	networkTimeout time.Duration = 2 * time.Minute
	pingInterval   time.Duration = time.Minute
	messageTimeout time.Duration = time.Minute

	sessionLimitBps int
	globalLimitBps  int
	sessionLimiter  *ratelimit.Bucket
	globalLimiter   *ratelimit.Bucket

	statusAddr       string
	poolAddrs        string
	defaultPoolAddrs string = "https://relays.syncthing.net/endpoint"
)

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)

	var dir, extAddress string

	flag.StringVar(&listen, "listen", ":22067", "Protocol listen address")
	flag.StringVar(&dir, "keys", ".", "Directory where cert.pem and key.pem is stored")
	flag.DurationVar(&networkTimeout, "network-timeout", networkTimeout, "Timeout for network operations between the client and the relay.\n\tIf no data is received between the client and the relay in this period of time, the connection is terminated.\n\tFurthermore, if no data is sent between either clients being relayed within this period of time, the session is also terminated.")
	flag.DurationVar(&pingInterval, "ping-interval", pingInterval, "How often pings are sent")
	flag.DurationVar(&messageTimeout, "message-timeout", messageTimeout, "Maximum amount of time we wait for relevant messages to arrive")
	flag.IntVar(&sessionLimitBps, "per-session-rate", sessionLimitBps, "Per session rate limit, in bytes/s")
	flag.IntVar(&globalLimitBps, "global-rate", globalLimitBps, "Global rate limit, in bytes/s")
	flag.BoolVar(&debug, "debug", debug, "Enable debug output")
	flag.StringVar(&statusAddr, "status-srv", ":22070", "Listen address for status service (blank to disable)")
	flag.StringVar(&poolAddrs, "pools", defaultPoolAddrs, "Comma separated list of relay pool addresses to join")

	flag.Parse()

	if extAddress == "" {
		extAddress = listen
	}

	addr, err := net.ResolveTCPAddr("tcp", extAddress)
	if err != nil {
		log.Fatal(err)
	}

	sessionAddress = addr.IP[:]
	sessionPort = uint16(addr.Port)

	certFile, keyFile := filepath.Join(dir, "cert.pem"), filepath.Join(dir, "key.pem")
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Println("Failed to load keypair. Generating one, this might take a while...")
		cert, err = tlsutil.NewCertificate(certFile, keyFile, "relaysrv", 3072)
		if err != nil {
			log.Fatalln("Failed to generate X509 key pair:", err)
		}
	}

	tlsCfg := &tls.Config{
		Certificates:           []tls.Certificate{cert},
		NextProtos:             []string{protocol.ProtocolName},
		ClientAuth:             tls.RequestClientCert,
		SessionTicketsDisabled: true,
		InsecureSkipVerify:     true,
		MinVersion:             tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
		},
	}

	id := syncthingprotocol.NewDeviceID(cert.Certificate[0])
	if debug {
		log.Println("ID:", id)
	}

	if sessionLimitBps > 0 {
		sessionLimiter = ratelimit.NewBucketWithRate(float64(sessionLimitBps), int64(2*sessionLimitBps))
	}
	if globalLimitBps > 0 {
		globalLimiter = ratelimit.NewBucketWithRate(float64(globalLimitBps), int64(2*globalLimitBps))
	}

	if statusAddr != "" {
		go statusService(statusAddr)
	}

	uri, err := url.Parse(fmt.Sprintf("relay://%s/?id=%s&pingInterval=%s&networkTimeout=%s&sessionLimitBps=%d&globalLimitBps=%d&statusAddr=%s", extAddress, id, pingInterval, networkTimeout, sessionLimitBps, globalLimitBps, statusAddr))
	if err != nil {
		log.Fatalln("Failed to construct URI", err)
	}

	log.Println("URI:", uri.String())

	if poolAddrs == defaultPoolAddrs {
		log.Println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
		log.Println("!!  Joining default relay pools, this relay will be available for public use. !!")
		log.Println(`!!      Use the -pools="" command line option to make the relay private.      !!`)
		log.Println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
	}

	pools := strings.Split(poolAddrs, ",")
	for _, pool := range pools {
		pool = strings.TrimSpace(pool)
		if len(pool) > 0 {
			go poolHandler(pool, uri)
		}
	}

	listener(listen, tlsCfg)
}
<commit_msg>Add ability to advertise provider<commit_after>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"net"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/juju/ratelimit"
	"github.com/syncthing/syncthing/lib/relay/protocol"
	"github.com/syncthing/syncthing/lib/tlsutil"

	syncthingprotocol "github.com/syncthing/syncthing/lib/protocol"
)

var (
	listen string
	debug  bool = false

	sessionAddress []byte
	sessionPort    uint16

	networkTimeout time.Duration = 2 * time.Minute
	pingInterval   time.Duration = time.Minute
	messageTimeout time.Duration = time.Minute

	sessionLimitBps int
	globalLimitBps  int
	sessionLimiter  *ratelimit.Bucket
	globalLimiter   *ratelimit.Bucket

	statusAddr       string
	poolAddrs        string
	providedBy       string
	defaultPoolAddrs string = "https://relays.syncthing.net/endpoint"
)

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)

	var dir, extAddress string

	flag.StringVar(&listen, "listen", ":22067", "Protocol listen address")
	flag.StringVar(&dir, "keys", ".", "Directory where cert.pem and key.pem is stored")
	flag.DurationVar(&networkTimeout, "network-timeout", networkTimeout, "Timeout for network operations between the client and the relay.\n\tIf no data is received between the client and the relay in this period of time, the connection is terminated.\n\tFurthermore, if no data is sent between either clients being relayed within this period of time, the session is also terminated.")
	flag.DurationVar(&pingInterval, "ping-interval", pingInterval, "How often pings are sent")
	flag.DurationVar(&messageTimeout, "message-timeout", messageTimeout, "Maximum amount of time we wait for relevant messages to arrive")
	flag.IntVar(&sessionLimitBps, "per-session-rate", sessionLimitBps, "Per session rate limit, in bytes/s")
	flag.IntVar(&globalLimitBps, "global-rate", globalLimitBps, "Global rate limit, in bytes/s")
	flag.BoolVar(&debug, "debug", debug, "Enable debug output")
	flag.StringVar(&statusAddr, "status-srv", ":22070", "Listen address for status service (blank to disable)")
	flag.StringVar(&poolAddrs, "pools", defaultPoolAddrs, "Comma separated list of relay pool addresses to join")
	flag.StringVar(&providedBy, "provided-by", "", "An optional description about who provides the relay")

	flag.Parse()

	if extAddress == "" {
		extAddress = listen
	}

	addr, err := net.ResolveTCPAddr("tcp", extAddress)
	if err != nil {
		log.Fatal(err)
	}

	sessionAddress = addr.IP[:]
	sessionPort = uint16(addr.Port)

	certFile, keyFile := filepath.Join(dir, "cert.pem"), filepath.Join(dir, "key.pem")
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Println("Failed to load keypair. Generating one, this might take a while...")
		cert, err = tlsutil.NewCertificate(certFile, keyFile, "relaysrv", 3072)
		if err != nil {
			log.Fatalln("Failed to generate X509 key pair:", err)
		}
	}

	tlsCfg := &tls.Config{
		Certificates:           []tls.Certificate{cert},
		NextProtos:             []string{protocol.ProtocolName},
		ClientAuth:             tls.RequestClientCert,
		SessionTicketsDisabled: true,
		InsecureSkipVerify:     true,
		MinVersion:             tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
		},
	}

	id := syncthingprotocol.NewDeviceID(cert.Certificate[0])
	if debug {
		log.Println("ID:", id)
	}

	if sessionLimitBps > 0 {
		sessionLimiter = ratelimit.NewBucketWithRate(float64(sessionLimitBps), int64(2*sessionLimitBps))
	}
	if globalLimitBps > 0 {
		globalLimiter = ratelimit.NewBucketWithRate(float64(globalLimitBps), int64(2*globalLimitBps))
	}

	if statusAddr != "" {
		go statusService(statusAddr)
	}

	uri, err := url.Parse(fmt.Sprintf("relay://%s/?id=%s&pingInterval=%s&networkTimeout=%s&sessionLimitBps=%d&globalLimitBps=%d&statusAddr=%s&providedBy=%s", extAddress, id, pingInterval, networkTimeout, sessionLimitBps, globalLimitBps, statusAddr, providedBy))
	if err != nil {
		log.Fatalln("Failed to construct URI", err)
	}

	log.Println("URI:", uri.String())

	if poolAddrs == defaultPoolAddrs {
		log.Println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
		log.Println("!!  Joining default relay pools, this relay will be available for public use. !!")
		log.Println(`!!      Use the -pools="" command line option to make the relay private.      !!`)
		log.Println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
	}

	pools := strings.Split(poolAddrs, ",")
	for _, pool := range pools {
		pool = strings.TrimSpace(pool)
		if len(pool) > 0 {
			go poolHandler(pool, uri)
		}
	}

	listener(listen, tlsCfg)
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Rook Authors. All rights reserved.

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

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

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

package ceph

import (
	"time"

	"github.com/rook/rook/cmd/rook/rook"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/daemon/ceph/client"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	cephver "github.com/rook/rook/pkg/operator/ceph/version"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/flags"
	"github.com/spf13/cobra"
)

var mgrCmd = &cobra.Command{
	Use: "mgr",
}
var mgrSidecarCmd = &cobra.Command{
	Use: "watch-active",
}
var (
	updateMgrServicesInterval string
	daemonName                string
	clusterSpec               cephv1.ClusterSpec
	rawCephVersion            string
)

func init() {
	addCephFlags(mgrCmd)

	// add the subcommands to the parent mgr command
	mgrCmd.AddCommand(mgrSidecarCmd)

	mgrSidecarCmd.Flags().BoolVar(&clusterSpec.Dashboard.Enabled, "dashboard-enabled", false, "whether the dashboard is enabled")
	mgrSidecarCmd.Flags().BoolVar(&clusterSpec.Monitoring.Enabled, "monitoring-enabled", false, "whether the monitoring is enabled")
	mgrSidecarCmd.Flags().StringVar(&updateMgrServicesInterval, "update-interval", "", "the interval at which to update the mgr services")
	mgrSidecarCmd.Flags().StringVar(&ownerRefID, "cluster-id", "", "the UID of the cluster CR that owns this cluster")
	mgrSidecarCmd.Flags().StringVar(&clusterName, "cluster-name", "", "the name of the cluster CR that owns this cluster")
	mgrSidecarCmd.Flags().StringVar(&daemonName, "daemon-name", "", "the name of the local mgr daemon")
	mgrSidecarCmd.Flags().StringVar(&rawCephVersion, "ceph-version", "", "the version of ceph")

	flags.SetFlagsFromEnv(mgrCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(mgrSidecarCmd.Flags(), rook.RookEnvVarPrefix)
	mgrSidecarCmd.RunE = runMgrSidecar
}

// Start the mgr daemon sidecar
func runMgrSidecar(cmd *cobra.Command, args []string) error {
	rook.SetLogLevel()

	context := createContext()
	clusterInfo.Monitors = mon.ParseMonEndpoints(cfg.monEndpoints)
	rook.LogStartupInfo(mgrSidecarCmd.Flags())

	ownerRef := opcontroller.ClusterOwnerRef(clusterName, ownerRefID)
	clusterInfo.OwnerInfo = k8sutil.NewOwnerInfoWithOwnerRef(&ownerRef, clusterInfo.Namespace)

	if err := client.WriteCephConfig(context, &clusterInfo); err != nil {
		rook.TerminateFatal(err)
	}

	interval, err := time.ParseDuration(updateMgrServicesInterval)
	if err != nil {
		rook.TerminateFatal(err)
	}

	version, err := cephver.ExtractCephVersion(rawCephVersion)
	if err != nil {
		rook.TerminateFatal(err)
	}
	clusterInfo.CephVersion = *version

	m := mgr.New(context, &clusterInfo, clusterSpec, "")
	for {
		err := m.ReconcileActiveMgrServices(daemonName)
		if err != nil {
			logger.Errorf("failed to reconcile services. %v", err)
		} else {
			logger.Infof("successfully reconciled services. checking again in %ds", (int)(interval.Seconds()))
		}
		time.Sleep(interval)
	}
}
<commit_msg>mgr: add missing context to the active-watch sidecar<commit_after>/*
Copyright 2021 The Rook Authors. All rights reserved.

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

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

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

package ceph

import (
	"time"

	"github.com/rook/rook/cmd/rook/rook"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/daemon/ceph/client"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	cephver "github.com/rook/rook/pkg/operator/ceph/version"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/flags"
	"github.com/spf13/cobra"
)

var mgrCmd = &cobra.Command{
	Use: "mgr",
}
var mgrSidecarCmd = &cobra.Command{
	Use: "watch-active",
}
var (
	updateMgrServicesInterval string
	daemonName                string
	clusterSpec               cephv1.ClusterSpec
	rawCephVersion            string
)

func init() {
	addCephFlags(mgrCmd)

	// add the subcommands to the parent mgr command
	mgrCmd.AddCommand(mgrSidecarCmd)

	mgrSidecarCmd.Flags().BoolVar(&clusterSpec.Dashboard.Enabled, "dashboard-enabled", false, "whether the dashboard is enabled")
	mgrSidecarCmd.Flags().BoolVar(&clusterSpec.Monitoring.Enabled, "monitoring-enabled", false, "whether the monitoring is enabled")
	mgrSidecarCmd.Flags().StringVar(&updateMgrServicesInterval, "update-interval", "", "the interval at which to update the mgr services")
	mgrSidecarCmd.Flags().StringVar(&ownerRefID, "cluster-id", "", "the UID of the cluster CR that owns this cluster")
	mgrSidecarCmd.Flags().StringVar(&clusterName, "cluster-name", "", "the name of the cluster CR that owns this cluster")
	mgrSidecarCmd.Flags().StringVar(&daemonName, "daemon-name", "", "the name of the local mgr daemon")
	mgrSidecarCmd.Flags().StringVar(&rawCephVersion, "ceph-version", "", "the version of ceph")

	flags.SetFlagsFromEnv(mgrCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(mgrSidecarCmd.Flags(), rook.RookEnvVarPrefix)
	mgrSidecarCmd.RunE = runMgrSidecar
}

// Start the mgr daemon sidecar
func runMgrSidecar(cmd *cobra.Command, args []string) error {
	rook.SetLogLevel()
	clusterInfo.Context = cmd.Context()

	context := createContext()
	clusterInfo.Monitors = mon.ParseMonEndpoints(cfg.monEndpoints)
	rook.LogStartupInfo(mgrSidecarCmd.Flags())

	ownerRef := opcontroller.ClusterOwnerRef(clusterName, ownerRefID)
	clusterInfo.OwnerInfo = k8sutil.NewOwnerInfoWithOwnerRef(&ownerRef, clusterInfo.Namespace)

	if err := client.WriteCephConfig(context, &clusterInfo); err != nil {
		rook.TerminateFatal(err)
	}

	interval, err := time.ParseDuration(updateMgrServicesInterval)
	if err != nil {
		rook.TerminateFatal(err)
	}

	version, err := cephver.ExtractCephVersion(rawCephVersion)
	if err != nil {
		rook.TerminateFatal(err)
	}
	clusterInfo.CephVersion = *version

	m := mgr.New(context, &clusterInfo, clusterSpec, "")
	for {
		err := m.ReconcileActiveMgrServices(daemonName)
		if err != nil {
			logger.Errorf("failed to reconcile services. %v", err)
		} else {
			logger.Infof("successfully reconciled services. checking again in %ds", (int)(interval.Seconds()))
		}
		time.Sleep(interval)
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/golang/glog"
)

func commandLine() {
	// treat this as a command line invocation, so read the input from the command line
	var request, response, engine, config string
	var help bool

	flag.BoolVar(&help, "h", false, "help")
	flag.StringVar(&request, "r", "", "Request file to upload (default: stdin)")
	flag.StringVar(&response, "s", "", "Response file (default: stdout)")
	flag.StringVar(&engine, "e", "", "Engine to use")
	flag.StringVar(&config, "c", "client_config.json", "Config file to read (json)")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: menagerie_cli [-v] [-r request file (json)] [-s response file (json)\n\tStandard input and output are used if the request/response are missing\n")
		flag.PrintDefaults()
		os.Exit(2)
	}
	flag.Parse()

	if help {
		flag.Usage()
		os.Exit(0)
	}

	var requestBody []byte
	var err error

	if request != "" {
		requestBody, err = ioutil.ReadFile(request)
		glog.V(0).Infoln("using request file %s", request)
	} else { // use stdin
		requestBody, err = ioutil.ReadAll(os.Stdin)
		glog.V(0).Infoln("using stdin")
	}

	if err != nil {
		glog.Fatalf("Can't read request file %s.", request)
	}

	var file = os.Stdout
	if response != "" {
		file, err = os.Create(response)
		if err != nil {
			glog.Fatal("Can't create response file. ", err)
		}
		defer file.Close()
	} else {
		glog.Infoln("No output file given. Using stdout")
	}

	glog.V(0).Infoln("responseFile %#v", file)

	// this event triggers it all
	t := NewClient(config, engine, file)
	_ = t.doit(requestBody)
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		flag.Usage()
	} else {
		commandLine()
	}

	os.Exit(0)
}
<commit_msg>better handling of cli flags<commit_after>package main

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

	"github.com/golang/glog"
)

func main() {
	// treat this as a command line invocation, so read the input from the command line
	var request, response, engine, config string
	var help bool

	flag.BoolVar(&help, "h", false, "help")
	flag.StringVar(&request, "r", "", "Request file to upload (default: stdin)")
	flag.StringVar(&response, "s", "", "Response file (default: stdout)")
	flag.StringVar(&engine, "e", "", "Engine to use")
	flag.StringVar(&config, "c", "client_config.json", "Config file to read (json)")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: menagerie_cli [-v] [-r request file (json)] [-s response file (json)\n\tStandard input and output are used if the request/response are missing\n")
		flag.PrintDefaults()
		os.Exit(2)
	}
	flag.Parse()

	if help || len(os.Args) == 1 {
		flag.Usage()
		os.Exit(0)
	}

	var requestBody []byte
	var err error

	if request != "" {
		requestBody, err = ioutil.ReadFile(request)
		glog.V(0).Infoln("using request file %s", request)
	} else { // use stdin
		requestBody, err = ioutil.ReadAll(os.Stdin)
		glog.V(0).Infoln("using stdin")
	}

	if err != nil {
		glog.Fatalf("Can't read request file %s.", request)
	}

	var file = os.Stdout
	if response != "" {
		file, err = os.Create(response)
		if err != nil {
			glog.Fatal("Can't create response file. ", err)
		}
		defer file.Close()
	} else {
		glog.Infoln("No output file given. Using stdout")
	}

	glog.V(0).Infoln("responseFile %#v", file)

	// this event triggers it all
	t := NewClient(config, engine, file)
	_ = t.doit(requestBody)
	os.Exit(0)
}
<|endoftext|>


<commit_before>package chat

// FIXME: Would be sweet if we could piggyback on a cli parser or something.

import (
	"bytes"
	"errors"
	"fmt"
	"strings"

	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/set"
)

// The error returned when an invalid command is issued.
var ErrInvalidCommand = errors.New("invalid command")

// The error returned when a command is given without an owner.
var ErrNoOwner = errors.New("command without owner")

// The error returned when a command is performed without the necessary number
// of arguments.
var ErrMissingArg = errors.New("missing argument")

// The error returned when a command is added without a prefix.
var ErrMissingPrefix = errors.New("command missing prefix")

// Command is a definition of a handler for a command.
type Command struct {
	// The command's key, such as /foo
	Prefix string
	// Extra help regarding arguments
	PrefixHelp string
	// If omitted, command is hidden from /help
	Help    string
	Handler func(*Room, message.CommandMsg) error
	// Command requires Op permissions
	Op bool
}

// Commands is a registry of available commands.
type Commands map[string]*Command

// Add will register a command. If help string is empty, it will be hidden from
// Help().
func (c Commands) Add(cmd Command) error {
	if cmd.Prefix == "" {
		return ErrMissingPrefix
	}

	c[cmd.Prefix] = &cmd
	return nil
}

// Alias will add another command for the same handler, won't get added to help.
func (c Commands) Alias(command string, alias string) error {
	cmd, ok := c[command]
	if !ok {
		return ErrInvalidCommand
	}
	c[alias] = cmd
	return nil
}

// Run executes a command message.
func (c Commands) Run(room *Room, msg message.CommandMsg) error {
	if msg.From() == nil {
		return ErrNoOwner
	}

	cmd, ok := c[msg.Command()]
	if !ok {
		return ErrInvalidCommand
	}

	return cmd.Handler(room, msg)
}

// Help will return collated help text as one string.
func (c Commands) Help(showOp bool) string {
	// Filter by op
	op := []*Command{}
	normal := []*Command{}
	for _, cmd := range c {
		if cmd.Op {
			op = append(op, cmd)
		} else {
			normal = append(normal, cmd)
		}
	}
	help := "Available commands:" + message.Newline + NewCommandsHelp(normal).String()
	if showOp {
		help += message.Newline + "-> Operator commands:" + message.Newline + NewCommandsHelp(op).String()
	}
	return help
}

var defaultCommands *Commands

func init() {
	defaultCommands = &Commands{}
	InitCommands(defaultCommands)
}

// InitCommands injects default commands into a Commands registry.
func InitCommands(c *Commands) {
	c.Add(Command{
		Prefix: "/help",
		Handler: func(room *Room, msg message.CommandMsg) error {
			op := room.IsOp(msg.From())
			room.Send(message.NewSystemMsg(room.commands.Help(op), msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/me",
		Handler: func(room *Room, msg message.CommandMsg) error {
			me := strings.TrimLeft(msg.Body(), "/me")
			if me == "" {
				me = "is at a loss for words."
			} else {
				me = me[1:]
			}

			room.Send(message.NewEmoteMsg(me, msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/exit",
		Help:   "Exit the chat.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			msg.From().Close()
			return nil
		},
	})
	c.Alias("/exit", "/quit")

	c.Add(Command{
		Prefix:     "/nick",
		PrefixHelp: "NAME",
		Help:       "Rename yourself.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			args := msg.Args()
			if len(args) != 1 {
				return ErrMissingArg
			}
			u := msg.From()

			member, ok := room.MemberByID(u.ID())
			if !ok {
				return errors.New("failed to find member")
			}

			oldID := member.ID()
			newID := SanitizeName(args[0])
			if newID == oldID {
				return errors.New("new name is the same as the original")
			}
			member.SetID(newID)
			err := room.Rename(oldID, member)
			if err != nil {
				member.SetID(oldID)
				return err
			}
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/names",
		Help:   "List users who are connected.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			theme := msg.From().Config().Theme

			colorize := func(u *message.User) string {
				return theme.ColorName(u)
			}

			if theme == nil {
				colorize = func(u *message.User) string {
					return u.Name()
				}
			}

			names := room.Members.ListPrefix("")
			colNames := make([]string, len(names))
			for i, uname := range names {
				colNames[i] = colorize(uname.Value().(*Member).User)
			}

			body := fmt.Sprintf("%d connected: %s", len(colNames), strings.Join(colNames, ", "))
			room.Send(message.NewSystemMsg(body, msg.From()))
			return nil
		},
	})
	c.Alias("/names", "/list")

	c.Add(Command{
		Prefix:     "/theme",
		PrefixHelp: "[colors|...]",
		Help:       "Set your color theme.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			user := msg.From()
			args := msg.Args()
			cfg := user.Config()
			if len(args) == 0 {
				theme := "plain"
				if cfg.Theme != nil {
					theme = cfg.Theme.ID()
				}
				var output bytes.Buffer
				fmt.Fprintf(&output, "Current theme: %s%s", theme, message.Newline)
				fmt.Fprintf(&output, "   Themes available: ")
				for i, t := range message.Themes {
					fmt.Fprintf(&output, t.ID())
					if i < len(message.Themes)-1 {
						fmt.Fprintf(&output, ", ")
					}
				}
				room.Send(message.NewSystemMsg(output.String(), user))
				return nil
			}

			id := args[0]
			for _, t := range message.Themes {
				if t.ID() == id {
					cfg.Theme = &t
					user.SetConfig(cfg)
					body := fmt.Sprintf("Set theme: %s", id)
					room.Send(message.NewSystemMsg(body, user))
					return nil
				}
			}
			return errors.New("theme not found")
		},
	})

	c.Add(Command{
		Prefix: "/quiet",
		Help:   "Silence room announcements.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			u := msg.From()
			cfg := u.Config()
			cfg.Quiet = !cfg.Quiet
			u.SetConfig(cfg)

			var body string
			if cfg.Quiet {
				body = "Quiet mode is toggled ON"
			} else {
				body = "Quiet mode is toggled OFF"
			}
			room.Send(message.NewSystemMsg(body, u))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/slap",
		PrefixHelp: "NAME",
		Handler: func(room *Room, msg message.CommandMsg) error {
			var me string
			args := msg.Args()
			if len(args) == 0 {
				me = "slaps themselves around a bit with a large trout."
			} else {
				me = fmt.Sprintf("slaps %s around a bit with a large trout.", strings.Join(args, " "))
			}

			room.Send(message.NewEmoteMsg(me, msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/timestamp",
		Help:   "Timestamps after 30min of inactivity.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			u := msg.From()
			cfg := u.Config()
			cfg.Timestamp = !cfg.Timestamp
			u.SetConfig(cfg)

			var body string
			if cfg.Timestamp {
				body = "Timestamp is toggled ON"
			} else {
				body = "Timestamp is toggled OFF"
			}
			room.Send(message.NewSystemMsg(body, u))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/ignore",
		PrefixHelp: "[USER]",
		Help:       "Hide messages from USER, /unignore USER to stop hiding.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/ignore"))
			if id == "" {
				// Print ignored names, if any.
				var names []string
				msg.From().Ignored.Each(func(_ string, item set.Item) error {
					names = append(names, item.Key())
					return nil
				})

				var systemMsg string
				if len(names) == 0 {
					systemMsg = "0 users ignored."
				} else {
					systemMsg = fmt.Sprintf("%d ignored: %s", len(names), strings.Join(names, ", "))
				}

				room.Send(message.NewSystemMsg(systemMsg, msg.From()))
				return nil
			}

			if id == msg.From().ID() {
				return errors.New("cannot ignore self")
			}
			target, ok := room.MemberByID(id)
			if !ok {
				return fmt.Errorf("user not found: %s", id)
			}

			err := msg.From().Ignored.Add(set.Itemize(id, target))
			if err == set.ErrCollision {
				return fmt.Errorf("user already ignored: %s", id)
			} else if err != nil {
				return err
			}

			room.Send(message.NewSystemMsg(fmt.Sprintf("Ignoring: %s", target.Name()), msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/unignore",
		PrefixHelp: "USER",
		Handler: func(room *Room, msg message.CommandMsg) error {
			id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/unignore"))
			if id == "" {
				return errors.New("must specify user")
			}

			if err := msg.From().Ignored.Remove(id); err != nil {
				return err
			}

			room.Send(message.NewSystemMsg(fmt.Sprintf("No longer ignoring: %s", id), msg.From()))
			return nil
		},
	})
}
<commit_msg>/shrug ¯\_(ツ)_/¯<commit_after>package chat

// FIXME: Would be sweet if we could piggyback on a cli parser or something.

import (
	"bytes"
	"errors"
	"fmt"
	"strings"

	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/set"
)

// The error returned when an invalid command is issued.
var ErrInvalidCommand = errors.New("invalid command")

// The error returned when a command is given without an owner.
var ErrNoOwner = errors.New("command without owner")

// The error returned when a command is performed without the necessary number
// of arguments.
var ErrMissingArg = errors.New("missing argument")

// The error returned when a command is added without a prefix.
var ErrMissingPrefix = errors.New("command missing prefix")

// Command is a definition of a handler for a command.
type Command struct {
	// The command's key, such as /foo
	Prefix string
	// Extra help regarding arguments
	PrefixHelp string
	// If omitted, command is hidden from /help
	Help    string
	Handler func(*Room, message.CommandMsg) error
	// Command requires Op permissions
	Op bool
}

// Commands is a registry of available commands.
type Commands map[string]*Command

// Add will register a command. If help string is empty, it will be hidden from
// Help().
func (c Commands) Add(cmd Command) error {
	if cmd.Prefix == "" {
		return ErrMissingPrefix
	}

	c[cmd.Prefix] = &cmd
	return nil
}

// Alias will add another command for the same handler, won't get added to help.
func (c Commands) Alias(command string, alias string) error {
	cmd, ok := c[command]
	if !ok {
		return ErrInvalidCommand
	}
	c[alias] = cmd
	return nil
}

// Run executes a command message.
func (c Commands) Run(room *Room, msg message.CommandMsg) error {
	if msg.From() == nil {
		return ErrNoOwner
	}

	cmd, ok := c[msg.Command()]
	if !ok {
		return ErrInvalidCommand
	}

	return cmd.Handler(room, msg)
}

// Help will return collated help text as one string.
func (c Commands) Help(showOp bool) string {
	// Filter by op
	op := []*Command{}
	normal := []*Command{}
	for _, cmd := range c {
		if cmd.Op {
			op = append(op, cmd)
		} else {
			normal = append(normal, cmd)
		}
	}
	help := "Available commands:" + message.Newline + NewCommandsHelp(normal).String()
	if showOp {
		help += message.Newline + "-> Operator commands:" + message.Newline + NewCommandsHelp(op).String()
	}
	return help
}

var defaultCommands *Commands

func init() {
	defaultCommands = &Commands{}
	InitCommands(defaultCommands)
}

// InitCommands injects default commands into a Commands registry.
func InitCommands(c *Commands) {
	c.Add(Command{
		Prefix: "/help",
		Handler: func(room *Room, msg message.CommandMsg) error {
			op := room.IsOp(msg.From())
			room.Send(message.NewSystemMsg(room.commands.Help(op), msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/me",
		Handler: func(room *Room, msg message.CommandMsg) error {
			me := strings.TrimLeft(msg.Body(), "/me")
			if me == "" {
				me = "is at a loss for words."
			} else {
				me = me[1:]
			}

			room.Send(message.NewEmoteMsg(me, msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/exit",
		Help:   "Exit the chat.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			msg.From().Close()
			return nil
		},
	})
	c.Alias("/exit", "/quit")

	c.Add(Command{
		Prefix:     "/nick",
		PrefixHelp: "NAME",
		Help:       "Rename yourself.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			args := msg.Args()
			if len(args) != 1 {
				return ErrMissingArg
			}
			u := msg.From()

			member, ok := room.MemberByID(u.ID())
			if !ok {
				return errors.New("failed to find member")
			}

			oldID := member.ID()
			newID := SanitizeName(args[0])
			if newID == oldID {
				return errors.New("new name is the same as the original")
			}
			member.SetID(newID)
			err := room.Rename(oldID, member)
			if err != nil {
				member.SetID(oldID)
				return err
			}
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/names",
		Help:   "List users who are connected.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			theme := msg.From().Config().Theme

			colorize := func(u *message.User) string {
				return theme.ColorName(u)
			}

			if theme == nil {
				colorize = func(u *message.User) string {
					return u.Name()
				}
			}

			names := room.Members.ListPrefix("")
			colNames := make([]string, len(names))
			for i, uname := range names {
				colNames[i] = colorize(uname.Value().(*Member).User)
			}

			body := fmt.Sprintf("%d connected: %s", len(colNames), strings.Join(colNames, ", "))
			room.Send(message.NewSystemMsg(body, msg.From()))
			return nil
		},
	})
	c.Alias("/names", "/list")

	c.Add(Command{
		Prefix:     "/theme",
		PrefixHelp: "[colors|...]",
		Help:       "Set your color theme.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			user := msg.From()
			args := msg.Args()
			cfg := user.Config()
			if len(args) == 0 {
				theme := "plain"
				if cfg.Theme != nil {
					theme = cfg.Theme.ID()
				}
				var output bytes.Buffer
				fmt.Fprintf(&output, "Current theme: %s%s", theme, message.Newline)
				fmt.Fprintf(&output, "   Themes available: ")
				for i, t := range message.Themes {
					fmt.Fprintf(&output, t.ID())
					if i < len(message.Themes)-1 {
						fmt.Fprintf(&output, ", ")
					}
				}
				room.Send(message.NewSystemMsg(output.String(), user))
				return nil
			}

			id := args[0]
			for _, t := range message.Themes {
				if t.ID() == id {
					cfg.Theme = &t
					user.SetConfig(cfg)
					body := fmt.Sprintf("Set theme: %s", id)
					room.Send(message.NewSystemMsg(body, user))
					return nil
				}
			}
			return errors.New("theme not found")
		},
	})

	c.Add(Command{
		Prefix: "/quiet",
		Help:   "Silence room announcements.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			u := msg.From()
			cfg := u.Config()
			cfg.Quiet = !cfg.Quiet
			u.SetConfig(cfg)

			var body string
			if cfg.Quiet {
				body = "Quiet mode is toggled ON"
			} else {
				body = "Quiet mode is toggled OFF"
			}
			room.Send(message.NewSystemMsg(body, u))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/slap",
		PrefixHelp: "NAME",
		Handler: func(room *Room, msg message.CommandMsg) error {
			var me string
			args := msg.Args()
			if len(args) == 0 {
				me = "slaps themselves around a bit with a large trout."
			} else {
				me = fmt.Sprintf("slaps %s around a bit with a large trout.", strings.Join(args, " "))
			}

			room.Send(message.NewEmoteMsg(me, msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/shrug",
		Handler: func(room *Room, msg message.CommandMsg) error {
			var me string
			args := msg.Args()
			if len(args) == 0 {
				me = `¯\_(ツ)_/¯`
			}

			room.Send(message.NewEmoteMsg(me, msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix: "/timestamp",
		Help:   "Timestamps after 30min of inactivity.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			u := msg.From()
			cfg := u.Config()
			cfg.Timestamp = !cfg.Timestamp
			u.SetConfig(cfg)

			var body string
			if cfg.Timestamp {
				body = "Timestamp is toggled ON"
			} else {
				body = "Timestamp is toggled OFF"
			}
			room.Send(message.NewSystemMsg(body, u))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/ignore",
		PrefixHelp: "[USER]",
		Help:       "Hide messages from USER, /unignore USER to stop hiding.",
		Handler: func(room *Room, msg message.CommandMsg) error {
			id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/ignore"))
			if id == "" {
				// Print ignored names, if any.
				var names []string
				msg.From().Ignored.Each(func(_ string, item set.Item) error {
					names = append(names, item.Key())
					return nil
				})

				var systemMsg string
				if len(names) == 0 {
					systemMsg = "0 users ignored."
				} else {
					systemMsg = fmt.Sprintf("%d ignored: %s", len(names), strings.Join(names, ", "))
				}

				room.Send(message.NewSystemMsg(systemMsg, msg.From()))
				return nil
			}

			if id == msg.From().ID() {
				return errors.New("cannot ignore self")
			}
			target, ok := room.MemberByID(id)
			if !ok {
				return fmt.Errorf("user not found: %s", id)
			}

			err := msg.From().Ignored.Add(set.Itemize(id, target))
			if err == set.ErrCollision {
				return fmt.Errorf("user already ignored: %s", id)
			} else if err != nil {
				return err
			}

			room.Send(message.NewSystemMsg(fmt.Sprintf("Ignoring: %s", target.Name()), msg.From()))
			return nil
		},
	})

	c.Add(Command{
		Prefix:     "/unignore",
		PrefixHelp: "USER",
		Handler: func(room *Room, msg message.CommandMsg) error {
			id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/unignore"))
			if id == "" {
				return errors.New("must specify user")
			}

			if err := msg.From().Ignored.Remove(id); err != nil {
				return err
			}

			room.Send(message.NewSystemMsg(fmt.Sprintf("No longer ignoring: %s", id), msg.From()))
			return nil
		},
	})
}
<|endoftext|>


<commit_before>package calls

import (
	"errors"
	"math/rand"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/scheduler"
)

// Filters sets a scheduler.Call's internal Filters, required for Accept and Decline calls.
func Filters(fo ...mesos.FilterOpt) scheduler.CallOpt {
	return func(c *scheduler.Call) {
		switch *c.Type {
		case scheduler.Call_ACCEPT:
			c.Accept.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_ACCEPT_INVERSE_OFFERS:
			c.AcceptInverseOffers.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_DECLINE:
			c.Decline.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_DECLINE_INVERSE_OFFERS:
			c.DeclineInverseOffers.Filters = mesos.OptionalFilters(fo...)
		default:
			panic("filters not supported for type " + c.Type.String())
		}
	}
}

// RefuseSecondsWithJitter returns a calls.Filters option that sets RefuseSeconds to a random number
// of seconds between 0 and the given duration.
func RefuseSecondsWithJitter(r *rand.Rand, d time.Duration) scheduler.CallOpt {
	return Filters(func(f *mesos.Filters) {
		s := time.Duration(r.Int63n(int64(d))).Seconds()
		f.RefuseSeconds = &s
	})
}

// Framework sets a scheduler.Call's FrameworkID
func Framework(id string) scheduler.CallOpt {
	return func(c *scheduler.Call) {
		c.FrameworkID = &mesos.FrameworkID{Value: id}
	}
}

// Subscribe returns a subscribe call with the given parameters.
// The call's FrameworkID is automatically filled in from the info specification.
func Subscribe(info *mesos.FrameworkInfo) *scheduler.Call {
	return &scheduler.Call{
		Type:        scheduler.Call_SUBSCRIBE.Enum(),
		FrameworkID: info.GetID(),
		Subscribe:   &scheduler.Call_Subscribe{FrameworkInfo: info},
	}
}

type acceptBuilder struct {
	offerIDs   map[mesos.OfferID]struct{}
	operations []mesos.Offer_Operation
}

type AcceptOpt func(*acceptBuilder)

type OfferOperations []mesos.Offer_Operation

// WithOffers allows a client to pair some set of OfferOperations with multiple resource offers.
// Example: calls.Accept(calls.OfferOperations{calls.OpLaunch(tasks...)}.WithOffers(offers...))
func (ob OfferOperations) WithOffers(ids ...mesos.OfferID) AcceptOpt {
	return func(ab *acceptBuilder) {
		for i := range ids {
			ab.offerIDs[ids[i]] = struct{}{}
		}
		ab.operations = append(ab.operations, ob...)
	}
}

// Accept returns an accept call with the given parameters.
// Callers are expected to fill in the FrameworkID and Filters.
func Accept(ops ...AcceptOpt) *scheduler.Call {
	ab := &acceptBuilder{
		offerIDs: make(map[mesos.OfferID]struct{}, len(ops)),
	}
	for _, op := range ops {
		op(ab)
	}
	offerIDs := make([]mesos.OfferID, 0, len(ab.offerIDs))
	for id := range ab.offerIDs {
		offerIDs = append(offerIDs, id)
	}
	return &scheduler.Call{
		Type: scheduler.Call_ACCEPT.Enum(),
		Accept: &scheduler.Call_Accept{
			OfferIDs:   offerIDs,
			Operations: ab.operations,
		},
	}
}

func AcceptInverseOffers(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_ACCEPT_INVERSE_OFFERS.Enum(),
		AcceptInverseOffers: &scheduler.Call_AcceptInverseOffers{
			InverseOfferIDs: offerIDs,
		},
	}
}

func DeclineInverseOffers(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_DECLINE_INVERSE_OFFERS.Enum(),
		DeclineInverseOffers: &scheduler.Call_DeclineInverseOffers{
			InverseOfferIDs: offerIDs,
		},
	}
}

// OpLaunch returns a launch operation builder for the given tasks
func OpLaunch(ti ...mesos.TaskInfo) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.LAUNCH.Enum(),
		Launch: &mesos.Offer_Operation_Launch{
			TaskInfos: ti,
		},
	}
}

func OpReserve(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.RESERVE.Enum(),
		Reserve: &mesos.Offer_Operation_Reserve{
			Resources: rs,
		},
	}
}

func OpUnreserve(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.UNRESERVE.Enum(),
		Unreserve: &mesos.Offer_Operation_Unreserve{
			Resources: rs,
		},
	}
}

func OpCreate(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.CREATE.Enum(),
		Create: &mesos.Offer_Operation_Create{
			Volumes: rs,
		},
	}
}

func OpDestroy(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.DESTROY.Enum(),
		Destroy: &mesos.Offer_Operation_Destroy{
			Volumes: rs,
		},
	}
}

// Revive returns a revive call.
// Callers are expected to fill in the FrameworkID.
func Revive() *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_REVIVE.Enum(),
	}
}

// Suppress returns a suppress call.
// Callers are expected to fill in the FrameworkID.
func Suppress() *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_SUPPRESS.Enum(),
	}
}

// Decline returns a decline call with the given parameters.
// Callers are expected to fill in the FrameworkID and Filters.
func Decline(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_DECLINE.Enum(),
		Decline: &scheduler.Call_Decline{
			OfferIDs: offerIDs,
		},
	}
}

// Kill returns a kill call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Kill(taskID, agentID string) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_KILL.Enum(),
		Kill: &scheduler.Call_Kill{
			TaskID:  mesos.TaskID{Value: taskID},
			AgentID: optionalAgentID(agentID),
		},
	}
}

// Shutdown returns a shutdown call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Shutdown(executorID, agentID string) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_SHUTDOWN.Enum(),
		Shutdown: &scheduler.Call_Shutdown{
			ExecutorID: mesos.ExecutorID{Value: executorID},
			AgentID:    mesos.AgentID{Value: agentID},
		},
	}
}

// Acknowledge returns an acknowledge call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Acknowledge(agentID, taskID string, uuid []byte) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_ACKNOWLEDGE.Enum(),
		Acknowledge: &scheduler.Call_Acknowledge{
			AgentID: mesos.AgentID{Value: agentID},
			TaskID:  mesos.TaskID{Value: taskID},
			UUID:    uuid,
		},
	}
}

// ReconcileTasks constructs a []Call_Reconcile_Task from the given mappings:
//     map[string]string{taskID:agentID}
// Map keys (taskID's) are required to be non-empty, but values (agentID's) *may* be empty.
func ReconcileTasks(tasks map[string]string) scheduler.ReconcileOpt {
	return func(cr *scheduler.Call_Reconcile) {
		if len(tasks) == 0 {
			cr.Tasks = nil
			return
		}
		result := make([]scheduler.Call_Reconcile_Task, len(tasks))
		i := 0
		for k, v := range tasks {
			result[i].TaskID = mesos.TaskID{Value: k}
			result[i].AgentID = optionalAgentID(v)
			i++
		}
		cr.Tasks = result
	}
}

// Reconcile returns a reconcile call with the given parameters.
// See ReconcileTask.
// Callers are expected to fill in the FrameworkID.
func Reconcile(opts ...scheduler.ReconcileOpt) *scheduler.Call {
	return &scheduler.Call{
		Type:      scheduler.Call_RECONCILE.Enum(),
		Reconcile: (&scheduler.Call_Reconcile{}).With(opts...),
	}
}

// Message returns a message call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Message(agentID, executorID string, data []byte) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_MESSAGE.Enum(),
		Message: &scheduler.Call_Message{
			AgentID:    mesos.AgentID{Value: agentID},
			ExecutorID: mesos.ExecutorID{Value: executorID},
			Data:       data,
		},
	}
}

// Request returns a resource request call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Request(requests ...mesos.Request) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_REQUEST.Enum(),
		Request: &scheduler.Call_Request{
			Requests: requests,
		},
	}
}

func optionalAgentID(agentID string) *mesos.AgentID {
	if agentID == "" {
		return nil
	}
	return &mesos.AgentID{Value: agentID}
}

func errInvalidCall(reason string) error {
	return errors.New("invalid call: " + reason)
}
<commit_msg>added docs for the new accept-inverse-offers and decline-inverse-offers calls<commit_after>package calls

import (
	"errors"
	"math/rand"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/scheduler"
)

// Filters sets a scheduler.Call's internal Filters, required for Accept and Decline calls.
func Filters(fo ...mesos.FilterOpt) scheduler.CallOpt {
	return func(c *scheduler.Call) {
		switch *c.Type {
		case scheduler.Call_ACCEPT:
			c.Accept.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_ACCEPT_INVERSE_OFFERS:
			c.AcceptInverseOffers.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_DECLINE:
			c.Decline.Filters = mesos.OptionalFilters(fo...)
		case scheduler.Call_DECLINE_INVERSE_OFFERS:
			c.DeclineInverseOffers.Filters = mesos.OptionalFilters(fo...)
		default:
			panic("filters not supported for type " + c.Type.String())
		}
	}
}

// RefuseSecondsWithJitter returns a calls.Filters option that sets RefuseSeconds to a random number
// of seconds between 0 and the given duration.
func RefuseSecondsWithJitter(r *rand.Rand, d time.Duration) scheduler.CallOpt {
	return Filters(func(f *mesos.Filters) {
		s := time.Duration(r.Int63n(int64(d))).Seconds()
		f.RefuseSeconds = &s
	})
}

// Framework sets a scheduler.Call's FrameworkID
func Framework(id string) scheduler.CallOpt {
	return func(c *scheduler.Call) {
		c.FrameworkID = &mesos.FrameworkID{Value: id}
	}
}

// Subscribe returns a subscribe call with the given parameters.
// The call's FrameworkID is automatically filled in from the info specification.
func Subscribe(info *mesos.FrameworkInfo) *scheduler.Call {
	return &scheduler.Call{
		Type:        scheduler.Call_SUBSCRIBE.Enum(),
		FrameworkID: info.GetID(),
		Subscribe:   &scheduler.Call_Subscribe{FrameworkInfo: info},
	}
}

type acceptBuilder struct {
	offerIDs   map[mesos.OfferID]struct{}
	operations []mesos.Offer_Operation
}

type AcceptOpt func(*acceptBuilder)

type OfferOperations []mesos.Offer_Operation

// WithOffers allows a client to pair some set of OfferOperations with multiple resource offers.
// Example: calls.Accept(calls.OfferOperations{calls.OpLaunch(tasks...)}.WithOffers(offers...))
func (ob OfferOperations) WithOffers(ids ...mesos.OfferID) AcceptOpt {
	return func(ab *acceptBuilder) {
		for i := range ids {
			ab.offerIDs[ids[i]] = struct{}{}
		}
		ab.operations = append(ab.operations, ob...)
	}
}

// Accept returns an accept call with the given parameters.
// Callers are expected to fill in the FrameworkID and Filters.
func Accept(ops ...AcceptOpt) *scheduler.Call {
	ab := &acceptBuilder{
		offerIDs: make(map[mesos.OfferID]struct{}, len(ops)),
	}
	for _, op := range ops {
		op(ab)
	}
	offerIDs := make([]mesos.OfferID, 0, len(ab.offerIDs))
	for id := range ab.offerIDs {
		offerIDs = append(offerIDs, id)
	}
	return &scheduler.Call{
		Type: scheduler.Call_ACCEPT.Enum(),
		Accept: &scheduler.Call_Accept{
			OfferIDs:   offerIDs,
			Operations: ab.operations,
		},
	}
}

// AcceptInverseOffers returns an accept-inverse-offers call for the given offer IDs.
// Callers are expected to fill in the FrameworkID and Filters.
func AcceptInverseOffers(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_ACCEPT_INVERSE_OFFERS.Enum(),
		AcceptInverseOffers: &scheduler.Call_AcceptInverseOffers{
			InverseOfferIDs: offerIDs,
		},
	}
}

// DeclineInverseOffers returns a decline-inverse-offers call for the given offer IDs.
// Callers are expected to fill in the FrameworkID and Filters.
func DeclineInverseOffers(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_DECLINE_INVERSE_OFFERS.Enum(),
		DeclineInverseOffers: &scheduler.Call_DeclineInverseOffers{
			InverseOfferIDs: offerIDs,
		},
	}
}

// OpLaunch returns a launch operation builder for the given tasks
func OpLaunch(ti ...mesos.TaskInfo) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.LAUNCH.Enum(),
		Launch: &mesos.Offer_Operation_Launch{
			TaskInfos: ti,
		},
	}
}

func OpReserve(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.RESERVE.Enum(),
		Reserve: &mesos.Offer_Operation_Reserve{
			Resources: rs,
		},
	}
}

func OpUnreserve(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.UNRESERVE.Enum(),
		Unreserve: &mesos.Offer_Operation_Unreserve{
			Resources: rs,
		},
	}
}

func OpCreate(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.CREATE.Enum(),
		Create: &mesos.Offer_Operation_Create{
			Volumes: rs,
		},
	}
}

func OpDestroy(rs ...mesos.Resource) mesos.Offer_Operation {
	return mesos.Offer_Operation{
		Type: mesos.DESTROY.Enum(),
		Destroy: &mesos.Offer_Operation_Destroy{
			Volumes: rs,
		},
	}
}

// Revive returns a revive call.
// Callers are expected to fill in the FrameworkID.
func Revive() *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_REVIVE.Enum(),
	}
}

// Suppress returns a suppress call.
// Callers are expected to fill in the FrameworkID.
func Suppress() *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_SUPPRESS.Enum(),
	}
}

// Decline returns a decline call with the given parameters.
// Callers are expected to fill in the FrameworkID and Filters.
func Decline(offerIDs ...mesos.OfferID) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_DECLINE.Enum(),
		Decline: &scheduler.Call_Decline{
			OfferIDs: offerIDs,
		},
	}
}

// Kill returns a kill call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Kill(taskID, agentID string) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_KILL.Enum(),
		Kill: &scheduler.Call_Kill{
			TaskID:  mesos.TaskID{Value: taskID},
			AgentID: optionalAgentID(agentID),
		},
	}
}

// Shutdown returns a shutdown call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Shutdown(executorID, agentID string) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_SHUTDOWN.Enum(),
		Shutdown: &scheduler.Call_Shutdown{
			ExecutorID: mesos.ExecutorID{Value: executorID},
			AgentID:    mesos.AgentID{Value: agentID},
		},
	}
}

// Acknowledge returns an acknowledge call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Acknowledge(agentID, taskID string, uuid []byte) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_ACKNOWLEDGE.Enum(),
		Acknowledge: &scheduler.Call_Acknowledge{
			AgentID: mesos.AgentID{Value: agentID},
			TaskID:  mesos.TaskID{Value: taskID},
			UUID:    uuid,
		},
	}
}

// ReconcileTasks constructs a []Call_Reconcile_Task from the given mappings:
//     map[string]string{taskID:agentID}
// Map keys (taskID's) are required to be non-empty, but values (agentID's) *may* be empty.
func ReconcileTasks(tasks map[string]string) scheduler.ReconcileOpt {
	return func(cr *scheduler.Call_Reconcile) {
		if len(tasks) == 0 {
			cr.Tasks = nil
			return
		}
		result := make([]scheduler.Call_Reconcile_Task, len(tasks))
		i := 0
		for k, v := range tasks {
			result[i].TaskID = mesos.TaskID{Value: k}
			result[i].AgentID = optionalAgentID(v)
			i++
		}
		cr.Tasks = result
	}
}

// Reconcile returns a reconcile call with the given parameters.
// See ReconcileTask.
// Callers are expected to fill in the FrameworkID.
func Reconcile(opts ...scheduler.ReconcileOpt) *scheduler.Call {
	return &scheduler.Call{
		Type:      scheduler.Call_RECONCILE.Enum(),
		Reconcile: (&scheduler.Call_Reconcile{}).With(opts...),
	}
}

// Message returns a message call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Message(agentID, executorID string, data []byte) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_MESSAGE.Enum(),
		Message: &scheduler.Call_Message{
			AgentID:    mesos.AgentID{Value: agentID},
			ExecutorID: mesos.ExecutorID{Value: executorID},
			Data:       data,
		},
	}
}

// Request returns a resource request call with the given parameters.
// Callers are expected to fill in the FrameworkID.
func Request(requests ...mesos.Request) *scheduler.Call {
	return &scheduler.Call{
		Type: scheduler.Call_REQUEST.Enum(),
		Request: &scheduler.Call_Request{
			Requests: requests,
		},
	}
}

func optionalAgentID(agentID string) *mesos.AgentID {
	if agentID == "" {
		return nil
	}
	return &mesos.AgentID{Value: agentID}
}

func errInvalidCall(reason string) error {
	return errors.New("invalid call: " + reason)
}
<|endoftext|>


<commit_before>package sheduler

import (
	"errors"
	"log"
	"os"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	uuid "github.com/satori/go.uuid"
	"github.com/stretchr/testify/assert"
)

type TestPublisher struct {
	ch chan SheduleData
}

func (p *TestPublisher) Publishing(data SheduleData) error {
	go func() {
		// TODO: catch the close channel
		p.ch <- data
	}()

	return nil
}

func TestScheduleBoltdbRabbitmq(t *testing.T) {
	os.RemoveAll("./_testdb.db") // delete the old file

	db, err := bolt.Open("./_testdb.db", 0600, &bolt.Options{
		Timeout: 1 * time.Second,
	})
	defer func() {
		os.RemoveAll("./_testdb.db")
	}()
	assert.NoError(t, err)
	ch := make(chan SheduleData)

	s := NewBoltDBSheduler(
		db,
		&TestPublisher{ch},
		[]byte("name"),
		time.Millisecond*10,
		log.New(os.Stdout, "[sheduler]", -1),
	)
	counter := 0
	go func() {
		for {
			select {
			case <-ch:
				counter++
			}
		}
	}()
	s.Start()

	for i := 0; i < 10; i++ {
		go func() {

			for i := 0; i < 100; i++ {
				time.Sleep(time.Millisecond * time.Duration(i/10))
				err := s.Shedule(uuid.NewV4(), []byte("data"), time.Now().Add(-time.Millisecond*10*time.Duration(i)))
				assert.NoError(t, err, "shedule")
			}
		}()
	}

	time.Sleep(time.Second)
	assert.Equal(t, counter, 1000)

	countTasks := 0
	db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("name"))
		bucket.ForEach(func(k, v []byte) error {
			countTasks++
			return nil
		})
		return nil
	})

	assert.Equal(t, counter, 1000)
	assert.Equal(t, countTasks, 0)
}

type TestFailPublisher struct {
	TestPublisher
}

func (p *TestFailPublisher) Publishing(data SheduleData) error {
	return errors.New("fail")
}

func TestSheduler_failpublish(t *testing.T) {
	os.RemoveAll("./_testdb.db") // delete the old file

	db, err := bolt.Open("./_testdb.db", 0600, &bolt.Options{
		Timeout: 1 * time.Second,
	})
	defer func() {
		os.RemoveAll("./_testdb.db")
	}()
	assert.NoError(t, err)
	ch := make(chan SheduleData)

	s := NewBoltDBSheduler(
		db,
		&TestFailPublisher{TestPublisher{ch}},
		[]byte("name"),
		time.Millisecond*10,
		log.New(os.Stdout, "[sheduler]", -1),
	)
	counter := 0
	go func() {
		for {
			select {
			case <-ch:
				counter++
			}
		}
	}()

	abort := s.Start()

	go func() {
		<-abort
		t.Log("abort worker")
	}()

	for i := 0; i < 10; i++ {
		go func() {

			for i := 0; i < 100; i++ {
				time.Sleep(time.Millisecond * time.Duration(i/10))
				err := s.Shedule(uuid.NewV4(), []byte("data"), time.Now().Add(-time.Millisecond*10*time.Duration(i)))
				assert.NoError(t, err, "shedule")
			}
		}()
	}

	time.Sleep(time.Second)

	// check

	countTasks := 0
	db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("name"))
		bucket.ForEach(func(k, v []byte) error {
			countTasks++
			return nil
		})
		return nil
	})

	assert.Equal(t, counter, 0)
	assert.Equal(t, countTasks, 1000)
}
<commit_msg>update tests (for travis)<commit_after>package sheduler

import (
	"errors"
	"log"
	"os"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	uuid "github.com/satori/go.uuid"
	"github.com/stretchr/testify/assert"
)

type TestPublisher struct {
	ch chan SheduleData
}

func (p *TestPublisher) Publishing(data SheduleData) error {
	go func() {
		// TODO: catch the close channel
		p.ch <- data
	}()

	return nil
}

func TestScheduleBoltdbRabbitmq(t *testing.T) {
	os.RemoveAll("./_testdb.db") // delete the old file

	db, err := bolt.Open("./_testdb.db", 0600, &bolt.Options{
		Timeout: 1 * time.Second,
	})
	defer func() {
		os.RemoveAll("./_testdb.db")
	}()
	assert.NoError(t, err)
	ch := make(chan SheduleData)

	s := NewBoltDBSheduler(
		db,
		&TestPublisher{ch},
		[]byte("name"),
		time.Millisecond*10,
		log.New(os.Stdout, "[sheduler]", -1),
	)
	counter := 0
	go func() {
		for {
			select {
			case <-ch:
				counter++
			}
		}
	}()
	s.Start()

	for i := 0; i < 10; i++ {
		go func() {

			for i := 0; i < 50; i++ {
				time.Sleep(time.Millisecond * time.Duration(i/10))
				err := s.Shedule(uuid.NewV4(), []byte("data"), time.Now().Add(-time.Millisecond*10*time.Duration(i)))
				assert.NoError(t, err, "shedule")
			}
		}()
	}

	time.Sleep(time.Second)

	countTasks := 0
	db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("name"))
		bucket.ForEach(func(k, v []byte) error {
			countTasks++
			return nil
		})
		return nil
	})

	assert.Equal(t, counter, 500)
	assert.Equal(t, countTasks, 0)
}

type TestFailPublisher struct {
	TestPublisher
}

func (p *TestFailPublisher) Publishing(data SheduleData) error {
	return errors.New("fail")
}

func TestSheduler_failpublish(t *testing.T) {
	os.RemoveAll("./_testdb.db") // delete the old file

	db, err := bolt.Open("./_testdb.db", 0600, &bolt.Options{
		Timeout: 1 * time.Second,
	})
	defer func() {
		os.RemoveAll("./_testdb.db")
	}()
	assert.NoError(t, err)
	ch := make(chan SheduleData)

	s := NewBoltDBSheduler(
		db,
		&TestFailPublisher{TestPublisher{ch}},
		[]byte("name"),
		time.Millisecond*10,
		log.New(os.Stdout, "[sheduler]", -1),
	)
	counter := 0
	go func() {
		for {
			select {
			case <-ch:
				counter++
			}
		}
	}()

	abort := s.Start()

	go func() {
		<-abort
		t.Log("abort worker")
	}()

	for i := 0; i < 10; i++ {
		go func() {

			for i := 0; i < 50; i++ {
				time.Sleep(time.Millisecond * time.Duration(i/10))
				err := s.Shedule(uuid.NewV4(), []byte("data"), time.Now().Add(-time.Millisecond*10*time.Duration(i)))
				assert.NoError(t, err, "shedule")
			}
		}()
	}

	time.Sleep(time.Second)

	// check

	countTasks := 0
	db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("name"))
		bucket.ForEach(func(k, v []byte) error {
			countTasks++
			return nil
		})
		return nil
	})

	assert.Equal(t, counter, 0)
	assert.Equal(t, countTasks, 500)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

var bars []*uiprogress.Bar

// runClient spawns threads for parallel uplink/downlink via TCP
func runClient(connectionType string, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"codePhrase": codePhrase,
		"connection": connectionType,
	})
	var wg sync.WaitGroup
	wg.Add(numberConnections)

	uiprogress.Start()
	bars = make([]*uiprogress.Bar, numberConnections)
	for id := 0; id < numberConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", serverAddress+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			logger.Debugf("telling relay: %s", connectionType+"."+codePhrase)

			sendMessage(connectionType+"."+Hash(codePhrase), connection)
			if connectionType == "s" { // this is a sender
				if id == 0 {
					fmt.Println("waiting for other to connect")
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				sendFile(id, connection, codePhrase)
			} else { // this is a receiver
				// receive file
				logger.Debug("receive file")
				fileName, fileIV, fileSalt, fileHash = receiveFile(id, connection, codePhrase)
			}

		}(id)
	}
	wg.Wait()

	if connectionType == "r" {
		catFile(fileName)
		encrypted, err := ioutil.ReadFile(fileName + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		decrypted, err := Decrypt(encrypted, codePhrase, fileSalt, fileIV)
		if err != nil {
			log.Error(err)
			return
		}
		ioutil.WriteFile(fileName, decrypted, 0644)
		os.Remove(fileName + ".encrypted")
		log.Debugf("\n\n\ndownloaded hash: %s", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: %s", fileHash)

		if fileHash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", fileName)
		} else {
			fmt.Printf("\nDownloaded %s!", fileName)
		}
	}
}

func catFile(fileNameToReceive string) {
	// cat the file
	os.Remove(fileNameToReceive)
	finished, err := os.Create(fileNameToReceive + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < numberConnections; id++ {
		fh, err := os.Open(fileNameToReceive + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	}

}

func receiveFile(id int, connection net.Conn, codePhrase string) (fileNameToReceive string, iv string, salt string, hashOfFile string) {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for file size")

	bufferFileSize := make([]byte, 10)
	connection.Read(bufferFileSize)
	fileSize, _ := strconv.ParseInt(strings.Trim(string(bufferFileSize), ":"), 10, 64)
	logger.Debugf("filesize: %d", fileSize)

	bufferFileName := make([]byte, 64)
	connection.Read(bufferFileName)
	fileNameToReceive = strings.Trim(string(bufferFileName), ":")
	logger.Debugf("fileName: %v", fileNameToReceive)

	ivHex := make([]byte, BUFFERSIZE)
	connection.Read(ivHex)
	iv = strings.Trim(string(ivHex), ":")
	logger.Debugf("iv: %s", iv)

	saltHex := make([]byte, BUFFERSIZE)
	connection.Read(saltHex)
	salt = strings.Trim(string(saltHex), ":")
	logger.Debugf("salt: %s", salt)

	hashOfFileBytes := make([]byte, BUFFERSIZE)
	connection.Read(hashOfFileBytes)
	hashOfFile = strings.Trim(string(hashOfFileBytes), ":")
	logger.Debugf("hashOfFile: %s", hashOfFile)

	os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	newFile, err := os.Create(fileNameToReceive + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	bars[id] = uiprogress.AddBar(int(fileSize)/1024 + 1).AppendCompleted().PrependElapsed()

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		bars[id].Incr()
		if (fileSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (fileSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-fileSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-fileSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		//Increment the counter
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return
}

func sendFile(id int, connection net.Conn, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	// // Open the file that needs to be send to the client
	// file, err := os.Open(fileName + ".encrypted")
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }
	// defer file.Close()
	// // Get the filename and filesize
	// fileInfo, err := file.Stat()
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }

	numChunks := math.Ceil(float64(len(fileBytes)) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(numberConnections)))

	bytesPerConnection := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == numberConnections {
		bytesPerConnection = int64(len(fileBytes)) - (numberConnections-1)*bytesPerConnection
	}
	fileSize := fillString(strconv.FormatInt(int64(bytesPerConnection), 10), 10)

	fileNameToSend := fillString(path.Base(fileName), 64)

	if id == 0 || id == numberConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerConnection: %v", bytesPerConnection)
		logger.Debugf("fileNameToSend: %v", path.Base(fileName))
	}

	// send file size
	logger.Debugf("sending fileSize: %s", fileSize)
	connection.Write([]byte(fileSize))

	// send fileName
	logger.Debugf("sending fileNameToSend: %s", fileNameToSend)
	connection.Write([]byte(fileNameToSend))

	// send iv
	logger.Debugf("sending iv: %s", fileIV)
	connection.Write([]byte(fillString(fileIV, BUFFERSIZE)))

	// send salt
	logger.Debugf("sending salt: %s", fileSalt)
	connection.Write([]byte(fillString(fileSalt, BUFFERSIZE)))

	// send sha256sum of file
	logger.Debugf("sending sha256sum: %s", fileHash)
	connection.Write([]byte(fillString(fileHash, BUFFERSIZE)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(fileBytes)
	chunkI := 0
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == numberConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<commit_msg>Add more trimming<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

var bars []*uiprogress.Bar

// runClient spawns threads for parallel uplink/downlink via TCP
func runClient(connectionType string, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"codePhrase": codePhrase,
		"connection": connectionType,
	})
	var wg sync.WaitGroup
	wg.Add(numberConnections)

	uiprogress.Start()
	bars = make([]*uiprogress.Bar, numberConnections)
	for id := 0; id < numberConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", serverAddress+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			logger.Debugf("telling relay: %s", connectionType+"."+codePhrase)

			sendMessage(connectionType+"."+Hash(codePhrase), connection)
			if connectionType == "s" { // this is a sender
				if id == 0 {
					fmt.Println("waiting for other to connect")
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				sendFile(id, connection, codePhrase)
			} else { // this is a receiver
				// receive file
				logger.Debug("receive file")
				fileName, fileIV, fileSalt, fileHash = receiveFile(id, connection, codePhrase)
			}

		}(id)
	}
	wg.Wait()

	if connectionType == "r" {
		catFile(fileName)
		encrypted, err := ioutil.ReadFile(fileName + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		decrypted, err := Decrypt(encrypted, codePhrase, fileSalt, fileIV)
		if err != nil {
			log.Error(err)
			return
		}
		ioutil.WriteFile(fileName, decrypted, 0644)
		os.Remove(fileName + ".encrypted")
		log.Debugf("\n\n\ndownloaded hash: %s", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: %s", fileHash)

		if fileHash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", fileName)
		} else {
			fmt.Printf("\nDownloaded %s!", fileName)
		}
	}
}

func catFile(fileNameToReceive string) {
	// cat the file
	os.Remove(fileNameToReceive)
	finished, err := os.Create(fileNameToReceive + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < numberConnections; id++ {
		fh, err := os.Open(fileNameToReceive + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	}

}

func receiveFile(id int, connection net.Conn, codePhrase string) (fileNameToReceive string, iv string, salt string, hashOfFile string) {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for file size")

	bufferFileSize := make([]byte, 10)
	connection.Read(bufferFileSize)
	fileSize, _ := strconv.ParseInt(strings.Trim(string(bufferFileSize), ":"), 10, 64)
	logger.Debugf("filesize: %d", fileSize)

	bufferFileName := make([]byte, 64)
	connection.Read(bufferFileName)
	fileNameToReceive = strings.Trim(string(bufferFileName), ":")
	logger.Debugf("fileName: %v", fileNameToReceive)

	ivHex := make([]byte, BUFFERSIZE)
	connection.Read(ivHex)
	iv = strings.Trim(strings.TrimSpace(string(ivHex)), ":")
	logger.Debugf("iv: %s", iv)

	saltHex := make([]byte, BUFFERSIZE)
	connection.Read(saltHex)
	salt = strings.Trim(strings.TrimSpace(string(saltHex)), ":")
	logger.Debugf("salt: %s", salt)

	hashOfFileBytes := make([]byte, BUFFERSIZE)
	connection.Read(hashOfFileBytes)
	hashOfFile = strings.Trim(strings.TrimSpace(string(hashOfFileBytes)), ":")
	logger.Debugf("hashOfFile: %s", hashOfFile)

	os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	newFile, err := os.Create(fileNameToReceive + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	bars[id] = uiprogress.AddBar(int(fileSize)/1024 + 1).AppendCompleted().PrependElapsed()

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		bars[id].Incr()
		if (fileSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (fileSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-fileSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-fileSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		//Increment the counter
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return
}

func sendFile(id int, connection net.Conn, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	// // Open the file that needs to be send to the client
	// file, err := os.Open(fileName + ".encrypted")
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }
	// defer file.Close()
	// // Get the filename and filesize
	// fileInfo, err := file.Stat()
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }

	numChunks := math.Ceil(float64(len(fileBytes)) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(numberConnections)))

	bytesPerConnection := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == numberConnections {
		bytesPerConnection = int64(len(fileBytes)) - (numberConnections-1)*bytesPerConnection
	}
	fileSize := fillString(strconv.FormatInt(int64(bytesPerConnection), 10), 10)

	fileNameToSend := fillString(path.Base(fileName), 64)

	if id == 0 || id == numberConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerConnection: %v", bytesPerConnection)
		logger.Debugf("fileNameToSend: %v", path.Base(fileName))
	}

	// send file size
	logger.Debugf("sending fileSize: %s", fileSize)
	connection.Write([]byte(fileSize))

	// send fileName
	logger.Debugf("sending fileNameToSend: %s", fileNameToSend)
	connection.Write([]byte(fileNameToSend))

	// send iv
	logger.Debugf("sending iv: %s", fileIV)
	connection.Write([]byte(fillString(fileIV, BUFFERSIZE)))

	// send salt
	logger.Debugf("sending salt: %s", fileSalt)
	connection.Write([]byte(fillString(fileSalt, BUFFERSIZE)))

	// send sha256sum of file
	logger.Debugf("sending sha256sum: %s", fileHash)
	connection.Write([]byte(fillString(fileHash, BUFFERSIZE)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(fileBytes)
	chunkI := 0
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == numberConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<|endoftext|>


<commit_before>package wpa_connect

import (
	"errors"

	"github.com/godbus/dbus"
	"github.com/mark2b/wpa-connect/log"
	"github.com/mark2b/wpa-connect/wpa_cli"
	"github.com/mark2b/wpa-connect/wpa_dbus"
)

type connectContext struct {
	phaseWaitForScanDone           bool
	phaseWaitForInterfaceConnected bool
	scanDone                       chan bool
	connectDone                    chan bool
}

func (self *wifiManager) Connect(ssid string, password string) (e error) {
	self.connectContext = &connectContext{}
	self.connectContext.scanDone = make(chan bool)
	self.connectContext.connectDone = make(chan bool)
	if wpa, err := wpa_dbus.NewWPA(); err == nil {
		wpa.WaifForSignals(self.OnSignal)
		wpa.AddSignalsObserver()
		if wpa.ReadInterface(self.NetInterface); wpa.Error == nil {
			iface := wpa.Interface
			iface.AddSignalsObserver()
			self.connectContext.phaseWaitForScanDone = true
			if iface.Scan(); wpa.Error == nil {
				// Wait for scan done
				<-self.connectContext.scanDone
				if iface.ReadBSSList(); iface.Error == nil {
					// Look for target BSS
					var bssFound = false
					for _, bss := range iface.BSSs {
						if bss.ReadSSID(); bss.Error == nil {
							log.Log.Info(bss.SSID)
							if ssid == bss.SSID {
								bssFound = true
								if err := self.ConnectToBSS(&bss, iface, password); err == nil {
									// Wait for connection
									cli := wpa_cli.WPACli{NetInterface: self.NetInterface}
									if err := cli.SaveConfig(); err == nil {
									} else {
										e = err
									}
								} else {
									e = err
								}
								break
							}
						} else {
							e = bss.Error
						}
					}
					if !bssFound {
						e = errors.New("ssid_not_found")
					}
				} else {
					e = iface.Error
				}
			} else {
				e = wpa.Error
			}
			iface.RemoveSignalsObserver()
		} else {
			e = wpa.Error
		}
		wpa.RemoveSignalsObserver()
		wpa.StopWaifForSignals()
	} else {
		e = err
	}
	log.Log.Debug("Connect exit")
	return
}

func (self *wifiManager) ConnectToBSS(bss *wpa_dbus.BSSWPA, iface *wpa_dbus.InterfaceWPA, password string) (e error) {
	addNetworkArgs := map[string]dbus.Variant{
		"ssid": dbus.MakeVariant(bss.SSID),
		"psk":  dbus.MakeVariant(password)}
	if iface.RemoveAllNetworks().AddNetwork(addNetworkArgs); iface.Error == nil {
		network := iface.NewNetwork
		self.connectContext.phaseWaitForInterfaceConnected = true
		if network.Select(); network.Error == nil {
			connected := <-self.connectContext.connectDone
			log.Log.Debug("Connected", connected)
			if !connected {
				e = errors.New("connection_failed")
			}
		} else {
			e = network.Error
		}
	} else {
		e = iface.Error
	}
	return
}

func (self *wifiManager) OnSignal(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug(signal.Name, signal.Path)
	switch signal.Name {
	case "fi.w1.wpa_supplicant1.Interface.BSSAdded":
	case "fi.w1.wpa_supplicant1.Interface.BSSRemoved":
		break
	case "fi.w1.wpa_supplicant1.Interface.ScanDone":
		self.processScanDone(wpa, signal)
	case "fi.w1.wpa_supplicant1.Interface.PropertiesChanged":
		log.Log.Debug(signal.Name, signal.Path, signal.Body)
		self.processInterfacePropertiesChanged(wpa, signal)
	default:
		log.Log.Debug(signal.Name, signal.Path, signal.Body)
	}
}

func (self *wifiManager) processScanDone(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug("processScanDone")
	if self.connectContext.phaseWaitForScanDone {
		self.connectContext.phaseWaitForScanDone = false
		self.connectContext.scanDone <- true
	}
}

func (self *wifiManager) processInterfacePropertiesChanged(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug("processInterfacePropertiesChanged")
	log.Log.Debug("phaseWaitForInterfaceConnected", self.connectContext.phaseWaitForInterfaceConnected)
	if self.connectContext.phaseWaitForInterfaceConnected {
		if len(signal.Body) > 0 {
			properties := signal.Body[0].(map[string]dbus.Variant)
			if stateVariant, hasState := properties["State"]; hasState {
				if state, ok := stateVariant.Value().(string); ok {
					log.Log.Debug("State", state)
					if state == "completed" {
						self.connectContext.phaseWaitForInterfaceConnected = false
						self.connectContext.connectDone <- true
						return
					} else if state == "disconnected" {
						self.connectContext.phaseWaitForInterfaceConnected = false
						self.connectContext.connectDone <- false
						return
					}
				}
			}
		}
	}
}

type wifiManager struct {
	connectContext *connectContext
	NetInterface   string
}
<commit_msg>no message<commit_after>package wpa_connect

import (
	"errors"

	"github.com/godbus/dbus"
	"github.com/mark2b/wpa-connect/log"
	"github.com/mark2b/wpa-connect/wpa_cli"
	"github.com/mark2b/wpa-connect/wpa_dbus"
)

type connectContext struct {
	phaseWaitForScanDone           bool
	phaseWaitForInterfaceConnected bool
	scanDone                       chan bool
	connectDone                    chan bool
}

func (self *wifiManager) Connect(ssid string, password string) (e error) {
	self.connectContext = &connectContext{}
	self.connectContext.scanDone = make(chan bool)
	self.connectContext.connectDone = make(chan bool)
	if wpa, err := wpa_dbus.NewWPA(); err == nil {
		wpa.WaifForSignals(self.OnSignal)
		wpa.AddSignalsObserver()
		if wpa.ReadInterface(self.NetInterface); wpa.Error == nil {
			iface := wpa.Interface
			iface.AddSignalsObserver()
			self.connectContext.phaseWaitForScanDone = true
			if iface.Scan(); wpa.Error == nil {
				// Wait for scan done
				<-self.connectContext.scanDone
				if iface.ReadBSSList(); iface.Error == nil {
					// Look for target BSS
					var bssFound = false
					for _, bss := range iface.BSSs {
						if bss.ReadSSID(); bss.Error == nil {
							log.Log.Info(bss.SSID)
							if ssid == bss.SSID {
								bssFound = true
								if err := self.ConnectToBSS(&bss, iface, password); err == nil {
									// Wait for connection
									cli := wpa_cli.WPACli{NetInterface: self.NetInterface}
									if err := cli.SaveConfig(); err == nil {
									} else {
										e = err
									}
								} else {
									e = err
								}
								break
							}
						} else {
							e = bss.Error
						}
					}
					if !bssFound {
						e = errors.New("ssid_not_found")
					}
				} else {
					e = iface.Error
				}
			} else {
				e = wpa.Error
			}
			iface.RemoveSignalsObserver()
		} else {
			e = wpa.Error
		}
		wpa.RemoveSignalsObserver()
		wpa.StopWaifForSignals()
	} else {
		e = err
	}
	log.Log.Debug("Connect exit")
	return
}

func (self *wifiManager) ConnectToBSS(bss *wpa_dbus.BSSWPA, iface *wpa_dbus.InterfaceWPA, password string) (e error) {
	addNetworkArgs := map[string]dbus.Variant{
		"ssid": dbus.MakeVariant(bss.SSID),
		"psk":  dbus.MakeVariant(password)}
	if iface.RemoveAllNetworks().AddNetwork(addNetworkArgs); iface.Error == nil {
		network := iface.NewNetwork
		self.connectContext.phaseWaitForInterfaceConnected = true
		if network.Select(); network.Error == nil {
			connected := <-self.connectContext.connectDone
			log.Log.Debug("Connected", connected)
			if !connected {
				e = errors.New("connection_failed")
			}
		} else {
			e = network.Error
		}
	} else {
		e = iface.Error
	}
	return
}

func (self *wifiManager) OnSignal(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug(signal.Name, signal.Path)
	switch signal.Name {
	case "fi.w1.wpa_supplicant1.Interface.BSSAdded":
	case "fi.w1.wpa_supplicant1.Interface.BSSRemoved":
		break
	case "fi.w1.wpa_supplicant1.Interface.ScanDone":
		self.processScanDone(wpa, signal)
	case "fi.w1.wpa_supplicant1.Interface.PropertiesChanged":
		log.Log.Debug(signal.Name, signal.Path, signal.Body)
		self.processInterfacePropertiesChanged(wpa, signal)
	default:
		log.Log.Debug(signal.Name, signal.Path, signal.Body)
	}
}

func (self *wifiManager) processScanDone(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug("processScanDone")
	if self.connectContext.phaseWaitForScanDone {
		self.connectContext.phaseWaitForScanDone = false
		self.connectContext.scanDone <- true
	}
}

func (self *wifiManager) processInterfacePropertiesChanged(wpa *wpa_dbus.WPA, signal *dbus.Signal) {
	log.Log.Debug("processInterfacePropertiesChanged")
	log.Log.Debug("phaseWaitForInterfaceConnected", self.connectContext.phaseWaitForInterfaceConnected)
	if self.connectContext.phaseWaitForInterfaceConnected {
		if len(signal.Body) > 0 {
			properties := signal.Body[0].(map[string]dbus.Variant)
			if stateVariant, hasState := properties["State"]; hasState {
				if state, ok := stateVariant.Value().(string); ok {
					log.Log.Debug("State", state)
					if state == "completed" {
						self.connectContext.phaseWaitForInterfaceConnected = false
						self.connectContext.connectDone <- true
						return
					} else if state == "disconnected" {
						self.connectContext.phaseWaitForInterfaceConnected = false
						self.connectContext.connectDone <- false
						return
					}
				}
			}
		}
	}
}

type wifiManager struct {
	connectContext *connectContext
	NetInterface   string
}

var (
	WifiManager = &wifiManager{NetInterface: "wlan0"}
)
<|endoftext|>


<commit_before>package sample

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"time"
)

// LessThan is a Sampler, which samples if the probe is lower than boundary, where boundary is calculated as
//     math.MaxUint64 / rate
type lessThan struct {
	sampleState
	boundary uint64
}

// NewLessThan returns a LessThan Sampler, random number generator is seeded with UnixNano
func NewLessThan(rate uint64) (Sampler, error) {
	return NewLessThanSeeded(rate, time.Now().UTC().UnixNano())
}

// NewLessThanSeeded returns a LessThan Sampler, allow for manual seeding of the random number generator
func NewLessThanSeeded(rate uint64, seed int64) (Sampler, error) {
	if rate == 0 {
		return nil, errors.New("rate must not be 0")
	}
	rnd := rand.New(rand.NewSource(seed))
	return &lessThan{sampleState: sampleState{rate: rate, seed: seed, rnd: rnd}, boundary: math.MaxUint64 / rate}, nil
}

func (lt *lessThan) Sample() bool {
	return lt.SampleFrom(randUint64(lt.rnd))
}

func (lt *lessThan) SampleFrom(probe uint64) bool {
	lt.sampleCount++
	if probe <= lt.boundary {
		lt.trueCount++
		return true
	}
	return false
}

func (lt *lessThan) String() string {
	type X *lessThan
	x := X(lt)
	return fmt.Sprintf("%+v", x)
}
<commit_msg>Fixed algorithm description<commit_after>package sample

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"time"
)

// LessThan is a Sampler, which samples if the probe is lower than boundary, where boundary is calculated as
//     math.MaxUint64 / rate
// The sampling is done by the following calculation:
//     probe <= lt.boundary
type lessThan struct {
	sampleState
	boundary uint64
}

// NewLessThan returns a LessThan Sampler, random number generator is seeded with UnixNano
func NewLessThan(rate uint64) (Sampler, error) {
	return NewLessThanSeeded(rate, time.Now().UTC().UnixNano())
}

// NewLessThanSeeded returns a LessThan Sampler, allow for manual seeding of the random number generator
func NewLessThanSeeded(rate uint64, seed int64) (Sampler, error) {
	if rate == 0 {
		return nil, errors.New("rate must not be 0")
	}
	rnd := rand.New(rand.NewSource(seed))
	return &lessThan{sampleState: sampleState{rate: rate, seed: seed, rnd: rnd}, boundary: math.MaxUint64 / rate}, nil
}

func (lt *lessThan) Sample() bool {
	return lt.SampleFrom(randUint64(lt.rnd))
}

func (lt *lessThan) SampleFrom(probe uint64) bool {
	lt.sampleCount++
	if probe <= lt.boundary {
		lt.trueCount++
		return true
	}
	return false
}

func (lt *lessThan) String() string {
	type X *lessThan
	x := X(lt)
	return fmt.Sprintf("%+v", x)
}
<|endoftext|>