<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package token implements a manager of serviceaccount tokens for pods running
// on the node.
package token

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	authenticationv1 "k8s.io/api/authentication/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/clock"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
)

const (
	maxTTL    = 24 * time.Hour
	gcPeriod  = time.Minute
	maxJitter = 10 * time.Second
)

// NewManager returns a new token manager.
func NewManager(c clientset.Interface) *Manager {
	// check whether the server supports token requests so we can give a more helpful error message
	supported := false
	once := &sync.Once{}
	tokenRequestsSupported := func() bool {
		once.Do(func() {
			resources, err := c.Discovery().ServerResourcesForGroupVersion("v1")
			if err != nil {
				return
			}
			for _, resource := range resources.APIResources {
				if resource.Name == "serviceaccounts/token" {
					supported = true
					return
				}
			}
		})
		return supported
	}

	m := &Manager{
		getToken: func(name, namespace string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error) {
			if c == nil {
				return nil, errors.New("cannot use TokenManager when kubelet is in standalone mode")
			}
			tokenRequest, err := c.CoreV1().ServiceAccounts(namespace).CreateToken(context.TODO(), name, tr, metav1.CreateOptions{})
			if apierrors.IsNotFound(err) && !tokenRequestsSupported() {
				return nil, fmt.Errorf("the API server does not have TokenRequest endpoints enabled")
			}
			return tokenRequest, err
		},
		cache: make(map[string]*authenticationv1.TokenRequest),
		clock: clock.RealClock{},
	}
	go wait.Forever(m.cleanup, gcPeriod)
	return m
}

// Manager manages service account tokens for pods.
type Manager struct {

	// cacheMutex guards the cache
	cacheMutex sync.RWMutex
	cache      map[string]*authenticationv1.TokenRequest

	// mocked for testing
	getToken func(name, namespace string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error)
	clock    clock.Clock
}

// GetServiceAccountToken gets a service account token for a pod from cache or
// from the TokenRequest API. This process is as follows:
// * Check the cache for the current token request.
// * If the token exists and does not require a refresh, return the current token.
// * Attempt to refresh the token.
// * If the token is refreshed successfully, save it in the cache and return the token.
// * If refresh fails and the old token is still valid, log an error and return the old token.
// * If refresh fails and the old token is no longer valid, return an error
func (m *Manager) GetServiceAccountToken(namespace, name string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error) {
	key := keyFunc(name, namespace, tr)

	ctr, ok := m.get(key)

	if ok && !m.requiresRefresh(ctr) {
		return ctr, nil
	}

	tr, err := m.getToken(name, namespace, tr)
	if err != nil {
		switch {
		case !ok:
			return nil, fmt.Errorf("failed to fetch token: %v", err)
		case m.expired(ctr):
			return nil, fmt.Errorf("token %s expired and refresh failed: %v", key, err)
		default:
			klog.Errorf("couldn't update token %s: %v", key, err)
			return ctr, nil
		}
	}

	m.set(key, tr)
	return tr, nil
}

// DeleteServiceAccountToken should be invoked when pod got deleted. It simply
// clean token manager cache.
func (m *Manager) DeleteServiceAccountToken(podUID types.UID) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	for k, tr := range m.cache {
		if tr.Spec.BoundObjectRef.UID == podUID {
			delete(m.cache, k)
		}
	}
}

func (m *Manager) cleanup() {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	for k, tr := range m.cache {
		if m.expired(tr) {
			delete(m.cache, k)
		}
	}
}

func (m *Manager) get(key string) (*authenticationv1.TokenRequest, bool) {
	m.cacheMutex.RLock()
	defer m.cacheMutex.RUnlock()
	ctr, ok := m.cache[key]
	return ctr, ok
}

func (m *Manager) set(key string, tr *authenticationv1.TokenRequest) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	m.cache[key] = tr
}

func (m *Manager) expired(t *authenticationv1.TokenRequest) bool {
	return m.clock.Now().After(t.Status.ExpirationTimestamp.Time)
}

// requiresRefresh returns true if the token is older than 80% of its total
// ttl, or if the token is older than 24 hours.
func (m *Manager) requiresRefresh(tr *authenticationv1.TokenRequest) bool {
	if tr.Spec.ExpirationSeconds == nil {
		cpy := tr.DeepCopy()
		cpy.Status.Token = ""
		klog.Errorf("expiration seconds was nil for tr: %#v", cpy)
		return false
	}
	now := m.clock.Now()
	exp := tr.Status.ExpirationTimestamp.Time
	iat := exp.Add(-1 * time.Duration(*tr.Spec.ExpirationSeconds) * time.Second)

	jitter := time.Duration(rand.Float64()*maxJitter.Seconds()) * time.Second
	if now.After(iat.Add(maxTTL - jitter)) {
		return true
	}
	// Require a refresh if within 20% of the TTL plus a jitter from the expiration time.
	if now.After(exp.Add(-1*time.Duration((*tr.Spec.ExpirationSeconds*20)/100)*time.Second - jitter)) {
		return true
	}
	return false
}

// keys should be nonconfidential and safe to log
func keyFunc(name, namespace string, tr *authenticationv1.TokenRequest) string {
	var exp int64
	if tr.Spec.ExpirationSeconds != nil {
		exp = *tr.Spec.ExpirationSeconds
	}

	var ref authenticationv1.BoundObjectReference
	if tr.Spec.BoundObjectRef != nil {
		ref = *tr.Spec.BoundObjectRef
	}

	return fmt.Sprintf("%q/%q/%#v/%#v/%#v", name, namespace, tr.Spec.Audiences, exp, ref)
}
<commit_msg>migrate pkg/kubelet/token to structured logging<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package token implements a manager of serviceaccount tokens for pods running
// on the node.
package token

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	authenticationv1 "k8s.io/api/authentication/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/clock"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
)

const (
	maxTTL    = 24 * time.Hour
	gcPeriod  = time.Minute
	maxJitter = 10 * time.Second
)

// NewManager returns a new token manager.
func NewManager(c clientset.Interface) *Manager {
	// check whether the server supports token requests so we can give a more helpful error message
	supported := false
	once := &sync.Once{}
	tokenRequestsSupported := func() bool {
		once.Do(func() {
			resources, err := c.Discovery().ServerResourcesForGroupVersion("v1")
			if err != nil {
				return
			}
			for _, resource := range resources.APIResources {
				if resource.Name == "serviceaccounts/token" {
					supported = true
					return
				}
			}
		})
		return supported
	}

	m := &Manager{
		getToken: func(name, namespace string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error) {
			if c == nil {
				return nil, errors.New("cannot use TokenManager when kubelet is in standalone mode")
			}
			tokenRequest, err := c.CoreV1().ServiceAccounts(namespace).CreateToken(context.TODO(), name, tr, metav1.CreateOptions{})
			if apierrors.IsNotFound(err) && !tokenRequestsSupported() {
				return nil, fmt.Errorf("the API server does not have TokenRequest endpoints enabled")
			}
			return tokenRequest, err
		},
		cache: make(map[string]*authenticationv1.TokenRequest),
		clock: clock.RealClock{},
	}
	go wait.Forever(m.cleanup, gcPeriod)
	return m
}

// Manager manages service account tokens for pods.
type Manager struct {

	// cacheMutex guards the cache
	cacheMutex sync.RWMutex
	cache      map[string]*authenticationv1.TokenRequest

	// mocked for testing
	getToken func(name, namespace string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error)
	clock    clock.Clock
}

// GetServiceAccountToken gets a service account token for a pod from cache or
// from the TokenRequest API. This process is as follows:
// * Check the cache for the current token request.
// * If the token exists and does not require a refresh, return the current token.
// * Attempt to refresh the token.
// * If the token is refreshed successfully, save it in the cache and return the token.
// * If refresh fails and the old token is still valid, log an error and return the old token.
// * If refresh fails and the old token is no longer valid, return an error
func (m *Manager) GetServiceAccountToken(namespace, name string, tr *authenticationv1.TokenRequest) (*authenticationv1.TokenRequest, error) {
	key := keyFunc(name, namespace, tr)

	ctr, ok := m.get(key)

	if ok && !m.requiresRefresh(ctr) {
		return ctr, nil
	}

	tr, err := m.getToken(name, namespace, tr)
	if err != nil {
		switch {
		case !ok:
			return nil, fmt.Errorf("failed to fetch token: %v", err)
		case m.expired(ctr):
			return nil, fmt.Errorf("token %s expired and refresh failed: %v", key, err)
		default:
			klog.ErrorS(err, "Couldn't update token", "cacheKey", key)
			return ctr, nil
		}
	}

	m.set(key, tr)
	return tr, nil
}

// DeleteServiceAccountToken should be invoked when pod got deleted. It simply
// clean token manager cache.
func (m *Manager) DeleteServiceAccountToken(podUID types.UID) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	for k, tr := range m.cache {
		if tr.Spec.BoundObjectRef.UID == podUID {
			delete(m.cache, k)
		}
	}
}

func (m *Manager) cleanup() {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	for k, tr := range m.cache {
		if m.expired(tr) {
			delete(m.cache, k)
		}
	}
}

func (m *Manager) get(key string) (*authenticationv1.TokenRequest, bool) {
	m.cacheMutex.RLock()
	defer m.cacheMutex.RUnlock()
	ctr, ok := m.cache[key]
	return ctr, ok
}

func (m *Manager) set(key string, tr *authenticationv1.TokenRequest) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	m.cache[key] = tr
}

func (m *Manager) expired(t *authenticationv1.TokenRequest) bool {
	return m.clock.Now().After(t.Status.ExpirationTimestamp.Time)
}

// requiresRefresh returns true if the token is older than 80% of its total
// ttl, or if the token is older than 24 hours.
func (m *Manager) requiresRefresh(tr *authenticationv1.TokenRequest) bool {
	if tr.Spec.ExpirationSeconds == nil {
		cpy := tr.DeepCopy()
		cpy.Status.Token = ""
		klog.ErrorS(nil, "Expiration seconds was nil for token request", "tokenRequest", cpy)
		return false
	}
	now := m.clock.Now()
	exp := tr.Status.ExpirationTimestamp.Time
	iat := exp.Add(-1 * time.Duration(*tr.Spec.ExpirationSeconds) * time.Second)

	jitter := time.Duration(rand.Float64()*maxJitter.Seconds()) * time.Second
	if now.After(iat.Add(maxTTL - jitter)) {
		return true
	}
	// Require a refresh if within 20% of the TTL plus a jitter from the expiration time.
	if now.After(exp.Add(-1*time.Duration((*tr.Spec.ExpirationSeconds*20)/100)*time.Second - jitter)) {
		return true
	}
	return false
}

// keys should be nonconfidential and safe to log
func keyFunc(name, namespace string, tr *authenticationv1.TokenRequest) string {
	var exp int64
	if tr.Spec.ExpirationSeconds != nil {
		exp = *tr.Spec.ExpirationSeconds
	}

	var ref authenticationv1.BoundObjectReference
	if tr.Spec.BoundObjectRef != nil {
		ref = *tr.Spec.BoundObjectRef
	}

	return fmt.Sprintf("%q/%q/%#v/%#v/%#v", name, namespace, tr.Spec.Audiences, exp, ref)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 command

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/kballard/go-shellquote"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/sync/errgroup"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/util"
)

var (
	layout = "2006-01-02 15:04:05.999999999 -0700"
)

// SSHRunner runs commands through SSH.
//
// It implements the CommandRunner interface.
type SSHRunner struct {
	c *ssh.Client
}

// NewSSHRunner returns a new SSHRunner that will run commands
// through the ssh.Client provided.
func NewSSHRunner(c *ssh.Client) *SSHRunner {
	return &SSHRunner{c}
}

// Remove runs a command to delete a file on the remote.
func (s *SSHRunner) Remove(f assets.CopyableFile) error {
	sess, err := s.c.NewSession()
	if err != nil {
		return errors.Wrap(err, "getting ssh session")
	}
	defer sess.Close()
	cmd := getDeleteFileCommand(f)
	return sess.Run(cmd)
}

// teeSSH runs an SSH command, streaming stdout, stderr to logs
func teeSSH(s *ssh.Session, cmd string, outB io.Writer, errB io.Writer) error {
	outPipe, err := s.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "stdout")
	}

	errPipe, err := s.StderrPipe()
	if err != nil {
		return errors.Wrap(err, "stderr")
	}
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		if err := teePrefix(util.ErrPrefix, errPipe, errB, glog.V(8).Infof); err != nil {
			glog.Errorf("tee stderr: %v", err)
		}
		wg.Done()
	}()
	go func() {
		if err := teePrefix(util.OutPrefix, outPipe, outB, glog.V(8).Infof); err != nil {
			glog.Errorf("tee stdout: %v", err)
		}
		wg.Done()
	}()
	err = s.Run(cmd)
	wg.Wait()
	return err
}

// RunCmd implements the Command Runner interface to run a exec.Cmd object
func (s *SSHRunner) RunCmd(cmd *exec.Cmd) (*RunResult, error) {
	rr := &RunResult{Args: cmd.Args}
	glog.Infof("(SSHRunner) Run:  %v", rr.Command())

	var outb, errb io.Writer
	start := time.Now()

	if cmd.Stdout == nil {
		var so bytes.Buffer
		outb = io.MultiWriter(&so, &rr.Stdout)
	} else {
		outb = io.MultiWriter(cmd.Stdout, &rr.Stdout)
	}

	if cmd.Stderr == nil {
		var se bytes.Buffer
		errb = io.MultiWriter(&se, &rr.Stderr)
	} else {
		errb = io.MultiWriter(cmd.Stderr, &rr.Stderr)
	}

	sess, err := s.c.NewSession()
	if err != nil {
		return rr, errors.Wrap(err, "NewSession")
	}

	defer func() {
		if err := sess.Close(); err != nil {
			if err != io.EOF {
				glog.Errorf("session close: %v", err)
			}
		}
	}()

	elapsed := time.Since(start)
	err = teeSSH(sess, shellquote.Join(cmd.Args...), outb, errb)
	if err == nil {
		// Reduce log spam
		if elapsed > (1 * time.Second) {
			glog.Infof("(SSHRunner) Done: %v: (%s)", rr.Command(), elapsed)
		}
	} else {
		if exitError, ok := err.(*exec.ExitError); ok {
			rr.ExitCode = exitError.ExitCode()
		}
		glog.Infof("(SSHRunner) Non-zero exit: %v: %v (%s)\n%s", rr.Command(), err, elapsed, rr.Output())
	}
	return rr, err
}

// Copy copies a file to the remote over SSH.
func (s *SSHRunner) Copy(f assets.CopyableFile) error {
	dst := path.Join(path.Join(f.GetTargetDir(), f.GetTargetName()))
	exists, err := s.sameFileExists(f, dst)
	if err != nil {
		glog.Infof("Checked if %s exists, but got error: %v", dst, err)
	}
	if exists {
		glog.Infof("Skipping copying %s as it already exists", dst)
		return nil
	}
	sess, err := s.c.NewSession()
	if err != nil {
		return errors.Wrap(err, "NewSession")
	}

	w, err := sess.StdinPipe()
	if err != nil {
		return errors.Wrap(err, "StdinPipe")
	}
	// The scpcmd below *should not* return until all data is copied and the
	// StdinPipe is closed. But let's use errgroup to make it explicit.
	var g errgroup.Group
	var copied int64
	glog.Infof("Transferring %d bytes to %s", f.GetLength(), dst)

	g.Go(func() error {
		defer w.Close()
		header := fmt.Sprintf("C%s %d %s\n", f.GetPermissions(), f.GetLength(), f.GetTargetName())
		fmt.Fprint(w, header)
		if f.GetLength() == 0 {
			glog.Warningf("%s is a 0 byte asset!", f.GetTargetName())
			fmt.Fprint(w, "\x00")
			return nil
		}

		copied, err = io.Copy(w, f)
		if err != nil {
			return errors.Wrap(err, "io.Copy")
		}
		if copied != int64(f.GetLength()) {
			return fmt.Errorf("%s: expected to copy %d bytes, but copied %d instead", f.GetTargetName(), f.GetLength(), copied)
		}
		glog.Infof("%s: copied %d bytes", f.GetTargetName(), copied)
		fmt.Fprint(w, "\x00")
		return nil
	})

	scp := fmt.Sprintf("sudo mkdir -p %s && sudo scp -t %s", f.GetTargetDir(), f.GetTargetDir())
	mtime, err := f.GetModTime()
	if err != nil {
		glog.Infof("error getting modtime for %s: %v", dst, err)
	} else {
		scp += fmt.Sprintf(" && sudo touch -d \"%s\" %s", mtime.Format(layout), dst)
	}
	out, err := sess.CombinedOutput(scp)
	if err != nil {
		return fmt.Errorf("%s: %s\noutput: %s", scp, err, out)
	}
	return g.Wait()
}

func (s *SSHRunner) sameFileExists(f assets.CopyableFile, dst string) (bool, error) {
	// get file size and modtime of the source
	srcSize := f.GetLength()
	srcModTime, err := f.GetModTime()
	if err != nil {
		return false, err
	}
	if srcModTime.IsZero() {
		return false, nil
	}

	// get file size and modtime of the destination
	sess, err := s.c.NewSession()
	if err != nil {
		return false, err
	}

	cmd := "stat -c \"%s %y\" " + dst
	out, err := sess.CombinedOutput(cmd)
	if err != nil {
		return false, err
	}
	outputs := strings.SplitN(strings.Trim(string(out), "\n"), " ", 2)

	dstSize, err := strconv.Atoi(outputs[0])
	if err != nil {
		return false, err
	}
	dstModTime, err := time.Parse(layout, outputs[1])
	if err != nil {
		return false, err
	}

	// compare sizes and modtimes
	if srcSize != dstSize {
		return false, errors.New("source file and destination file are different sizes")
	}
	return srcModTime.Equal(dstModTime), nil
}

// teePrefix copies bytes from a reader to writer, logging each new line.
func teePrefix(prefix string, r io.Reader, w io.Writer, logger func(format string, args ...interface{})) error {
	scanner := bufio.NewScanner(r)
	scanner.Split(bufio.ScanBytes)
	var line bytes.Buffer

	for scanner.Scan() {
		b := scanner.Bytes()
		if _, err := w.Write(b); err != nil {
			return err
		}
		if bytes.IndexAny(b, "\r\n") == 0 {
			if line.Len() > 0 {
				logger("%s%s", prefix, line.String())
				line.Reset()
			}
			continue
		}
		line.Write(b)
	}
	// Catch trailing output in case stream does not end with a newline
	if line.Len() > 0 {
		logger("%s%s", prefix, line.String())
	}
	return nil
}
<commit_msg>ssh: Log and return stdout/stderr consistency<commit_after>/*
Copyright 2016 The Kubernetes 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 command

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/kballard/go-shellquote"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/sync/errgroup"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/util"
)

var (
	layout = "2006-01-02 15:04:05.999999999 -0700"
)

// SSHRunner runs commands through SSH.
//
// It implements the CommandRunner interface.
type SSHRunner struct {
	c *ssh.Client
}

// NewSSHRunner returns a new SSHRunner that will run commands
// through the ssh.Client provided.
func NewSSHRunner(c *ssh.Client) *SSHRunner {
	return &SSHRunner{c}
}

// Remove runs a command to delete a file on the remote.
func (s *SSHRunner) Remove(f assets.CopyableFile) error {
	sess, err := s.c.NewSession()
	if err != nil {
		return errors.Wrap(err, "getting ssh session")
	}
	defer sess.Close()
	cmd := getDeleteFileCommand(f)
	return sess.Run(cmd)
}

// teeSSH runs an SSH command, streaming stdout, stderr to logs
func teeSSH(s *ssh.Session, cmd string, outB io.Writer, errB io.Writer) error {
	outPipe, err := s.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "stdout")
	}

	errPipe, err := s.StderrPipe()
	if err != nil {
		return errors.Wrap(err, "stderr")
	}
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		if err := teePrefix(util.ErrPrefix, errPipe, errB, glog.V(8).Infof); err != nil {
			glog.Errorf("tee stderr: %v", err)
		}
		wg.Done()
	}()
	go func() {
		if err := teePrefix(util.OutPrefix, outPipe, outB, glog.V(8).Infof); err != nil {
			glog.Errorf("tee stdout: %v", err)
		}
		wg.Done()
	}()
	err = s.Run(cmd)
	wg.Wait()
	return err
}

// RunCmd implements the Command Runner interface to run a exec.Cmd object
func (s *SSHRunner) RunCmd(cmd *exec.Cmd) (*RunResult, error) {
	rr := &RunResult{Args: cmd.Args}
	glog.Infof("Run: %v", rr.Command())

	var outb, errb io.Writer
	start := time.Now()

	if cmd.Stdout == nil {
		var so bytes.Buffer
		outb = io.MultiWriter(&so, &rr.Stdout)
	} else {
		outb = io.MultiWriter(cmd.Stdout, &rr.Stdout)
	}

	if cmd.Stderr == nil {
		var se bytes.Buffer
		errb = io.MultiWriter(&se, &rr.Stderr)
	} else {
		errb = io.MultiWriter(cmd.Stderr, &rr.Stderr)
	}

	sess, err := s.c.NewSession()
	if err != nil {
		return rr, errors.Wrap(err, "NewSession")
	}

	defer func() {
		if err := sess.Close(); err != nil {
			if err != io.EOF {
				glog.Errorf("session close: %v", err)
			}
		}
	}()

	err = teeSSH(sess, shellquote.Join(cmd.Args...), outb, errb)
	elapsed := time.Since(start)

	if exitError, ok := err.(*exec.ExitError); ok {
		rr.ExitCode = exitError.ExitCode()
	}
	// Decrease log spam
	if elapsed > (1 * time.Second) {
		glog.Infof("Exit %d for %v: (%s)", rr.Command(), elapsed)
	}
	if err == nil {
		return rr, nil
	}

	return rr, fmt.Errorf("%s: %v\nstdout: %s\nstderr: %s", rr.Command(), err, rr.Stdout.String(), rr.Stderr.String())
}

// Copy copies a file to the remote over SSH.
func (s *SSHRunner) Copy(f assets.CopyableFile) error {
	dst := path.Join(path.Join(f.GetTargetDir(), f.GetTargetName()))
	exists, err := s.sameFileExists(f, dst)
	if err != nil {
		glog.Infof("Checked if %s exists, but got error: %v", dst, err)
	}
	if exists {
		glog.Infof("Skipping copying %s as it already exists", dst)
		return nil
	}
	sess, err := s.c.NewSession()
	if err != nil {
		return errors.Wrap(err, "NewSession")
	}

	w, err := sess.StdinPipe()
	if err != nil {
		return errors.Wrap(err, "StdinPipe")
	}
	// The scpcmd below *should not* return until all data is copied and the
	// StdinPipe is closed. But let's use errgroup to make it explicit.
	var g errgroup.Group
	var copied int64
	glog.Infof("Transferring %d bytes to %s", f.GetLength(), dst)

	g.Go(func() error {
		defer w.Close()
		header := fmt.Sprintf("C%s %d %s\n", f.GetPermissions(), f.GetLength(), f.GetTargetName())
		fmt.Fprint(w, header)
		if f.GetLength() == 0 {
			glog.Warningf("%s is a 0 byte asset!", f.GetTargetName())
			fmt.Fprint(w, "\x00")
			return nil
		}

		copied, err = io.Copy(w, f)
		if err != nil {
			return errors.Wrap(err, "io.Copy")
		}
		if copied != int64(f.GetLength()) {
			return fmt.Errorf("%s: expected to copy %d bytes, but copied %d instead", f.GetTargetName(), f.GetLength(), copied)
		}
		glog.Infof("%s: copied %d bytes", f.GetTargetName(), copied)
		fmt.Fprint(w, "\x00")
		return nil
	})

	scp := fmt.Sprintf("sudo mkdir -p %s && sudo scp -t %s", f.GetTargetDir(), f.GetTargetDir())
	mtime, err := f.GetModTime()
	if err != nil {
		glog.Infof("error getting modtime for %s: %v", dst, err)
	} else {
		scp += fmt.Sprintf(" && sudo touch -d \"%s\" %s", mtime.Format(layout), dst)
	}
	out, err := sess.CombinedOutput(scp)
	if err != nil {
		return fmt.Errorf("%s: %s\noutput: %s", scp, err, out)
	}
	return g.Wait()
}

func (s *SSHRunner) sameFileExists(f assets.CopyableFile, dst string) (bool, error) {
	// get file size and modtime of the source
	srcSize := f.GetLength()
	srcModTime, err := f.GetModTime()
	if err != nil {
		return false, err
	}
	if srcModTime.IsZero() {
		return false, nil
	}

	// get file size and modtime of the destination
	sess, err := s.c.NewSession()
	if err != nil {
		return false, err
	}

	cmd := "stat -c \"%s %y\" " + dst
	out, err := sess.CombinedOutput(cmd)
	if err != nil {
		return false, err
	}
	outputs := strings.SplitN(strings.Trim(string(out), "\n"), " ", 2)

	dstSize, err := strconv.Atoi(outputs[0])
	if err != nil {
		return false, err
	}
	dstModTime, err := time.Parse(layout, outputs[1])
	if err != nil {
		return false, err
	}

	// compare sizes and modtimes
	if srcSize != dstSize {
		return false, errors.New("source file and destination file are different sizes")
	}
	return srcModTime.Equal(dstModTime), nil
}

// teePrefix copies bytes from a reader to writer, logging each new line.
func teePrefix(prefix string, r io.Reader, w io.Writer, logger func(format string, args ...interface{})) error {
	scanner := bufio.NewScanner(r)
	scanner.Split(bufio.ScanBytes)
	var line bytes.Buffer

	for scanner.Scan() {
		b := scanner.Bytes()
		if _, err := w.Write(b); err != nil {
			return err
		}
		if bytes.IndexAny(b, "\r\n") == 0 {
			if line.Len() > 0 {
				logger("%s%s", prefix, line.String())
				line.Reset()
			}
			continue
		}
		line.Write(b)
	}
	// Catch trailing output in case stream does not end with a newline
	if line.Len() > 0 {
		logger("%s%s", prefix, line.String())
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package healthz

import (
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

func TestInstallHandler(t *testing.T) {
	mux := http.NewServeMux()
	InstallHandler(mux)
	req, err := http.NewRequest("GET", "http://example.com/healthz", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}
}

func TestInstallPathHandler(t *testing.T) {
	mux := http.NewServeMux()
	InstallPathHandler(mux, "/healthz/test")
	InstallPathHandler(mux, "/healthz/ready")
	req, err := http.NewRequest("GET", "http://example.com/healthz/test", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}

	req, err = http.NewRequest("GET", "http://example.com/healthz/ready", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w = httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}

}

func TestMulitipleChecks(t *testing.T) {
	tests := []struct {
		path             string
		expectedResponse string
		expectedStatus   int
		addBadCheck      bool
	}{
		{"/healthz?verbose", "[+]ping ok\nhealthz check passed\n", http.StatusOK, false},
		{"/healthz/ping", "ok", http.StatusOK, false},
		{"/healthz", "ok", http.StatusOK, false},
		{"/healthz?verbose", "[+]ping ok\n[-]bad failed: reason withheld\nhealthz check failed\n", http.StatusInternalServerError, true},
		{"/healthz/ping", "ok", http.StatusOK, true},
		{"/healthz/bad", "internal server error: this will fail\n", http.StatusInternalServerError, true},
		{"/healthz", "[+]ping ok\n[-]bad failed: reason withheld\nhealthz check failed\n", http.StatusInternalServerError, true},
	}

	for i, test := range tests {
		mux := http.NewServeMux()
		checks := []HealthzChecker{PingHealthz}
		if test.addBadCheck {
			checks = append(checks, NamedCheck("bad", func(_ *http.Request) error {
				return errors.New("this will fail")
			}))
		}
		InstallHandler(mux, checks...)
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com%v", test.path), nil)
		if err != nil {
			t.Fatalf("case[%d] Unexpected error: %v", i, err)
		}
		w := httptest.NewRecorder()
		mux.ServeHTTP(w, req)
		if w.Code != test.expectedStatus {
			t.Errorf("case[%d] Expected: %v, got: %v", i, test.expectedStatus, w.Code)
		}
		if w.Body.String() != test.expectedResponse {
			t.Errorf("case[%d] Expected:\n%v\ngot:\n%v\n", i, test.expectedResponse, w.Body.String())
		}
	}
}

func TestCheckerNames(t *testing.T) {
	n1 := "n1"
	n2 := "n2"
	c1 := &healthzCheck{name: n1}
	c2 := &healthzCheck{name: n2}

	testCases := []struct {
		desc string
		have []HealthzChecker
		want []string
	}{
		{"no checker", []HealthzChecker{}, []string{}},
		{"one checker", []HealthzChecker{c1}, []string{n1}},
		{"other checker", []HealthzChecker{c2}, []string{n2}},
		{"checker order", []HealthzChecker{c1, c2}, []string{n1, n2}},
		{"different checker order", []HealthzChecker{c2, c1}, []string{n2, n1}},
	}

	for _, tc := range testCases {
		result := checkerNames(tc.have...)
		t.Run(tc.desc, func(t *testing.T) {
			if reflect.DeepEqual(tc.want, result) {
				t.Errorf("want %#v, got %#v", tc.want, result)
			}
		})
	}

}
<commit_msg>Improve unit tests for InstallPathHandler<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package healthz

import (
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

func TestInstallHandler(t *testing.T) {
	mux := http.NewServeMux()
	InstallHandler(mux)
	req, err := http.NewRequest("GET", "http://example.com/healthz", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}
}

func TestInstallPathHandler(t *testing.T) {
	mux := http.NewServeMux()
	InstallPathHandler(mux, "/healthz/test")
	InstallPathHandler(mux, "/healthz/ready")
	req, err := http.NewRequest("GET", "http://example.com/healthz/test", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}

	req, err = http.NewRequest("GET", "http://example.com/healthz/ready", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	w = httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("expected %v, got %v", http.StatusOK, w.Code)
	}
	if w.Body.String() != "ok" {
		t.Errorf("expected %v, got %v", "ok", w.Body.String())
	}

}

func testMultipleChecks(path string, t *testing.T) {
	tests := []struct {
		path             string
		expectedResponse string
		expectedStatus   int
		addBadCheck      bool
	}{
		{"?verbose", "[+]ping ok\nhealthz check passed\n", http.StatusOK, false},
		{"/ping", "ok", http.StatusOK, false},
		{"", "ok", http.StatusOK, false},
		{"?verbose", "[+]ping ok\n[-]bad failed: reason withheld\nhealthz check failed\n", http.StatusInternalServerError, true},
		{"/ping", "ok", http.StatusOK, true},
		{"/bad", "internal server error: this will fail\n", http.StatusInternalServerError, true},
		{"", "[+]ping ok\n[-]bad failed: reason withheld\nhealthz check failed\n", http.StatusInternalServerError, true},
	}

	for i, test := range tests {
		mux := http.NewServeMux()
		checks := []HealthzChecker{PingHealthz}
		if test.addBadCheck {
			checks = append(checks, NamedCheck("bad", func(_ *http.Request) error {
				return errors.New("this will fail")
			}))
		}
		if path == "" {
			InstallHandler(mux, checks...)
			path = "/healthz"
		} else {
			InstallPathHandler(mux, path, checks...)
		}
		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com%s%v", path, test.path), nil)
		if err != nil {
			t.Fatalf("case[%d] Unexpected error: %v", i, err)
		}
		w := httptest.NewRecorder()
		mux.ServeHTTP(w, req)
		if w.Code != test.expectedStatus {
			t.Errorf("case[%d] Expected: %v, got: %v", i, test.expectedStatus, w.Code)
		}
		if w.Body.String() != test.expectedResponse {
			t.Errorf("case[%d] Expected:\n%v\ngot:\n%v\n", i, test.expectedResponse, w.Body.String())
		}
	}
}

func TestMultipleChecks(t *testing.T) {
	testMultipleChecks("", t)
}

func TestMultiplePathChecks(t *testing.T) {
	testMultipleChecks("/ready", t)
}

func TestCheckerNames(t *testing.T) {
	n1 := "n1"
	n2 := "n2"
	c1 := &healthzCheck{name: n1}
	c2 := &healthzCheck{name: n2}

	testCases := []struct {
		desc string
		have []HealthzChecker
		want []string
	}{
		{"no checker", []HealthzChecker{}, []string{}},
		{"one checker", []HealthzChecker{c1}, []string{n1}},
		{"other checker", []HealthzChecker{c2}, []string{n2}},
		{"checker order", []HealthzChecker{c1, c2}, []string{n1, n2}},
		{"different checker order", []HealthzChecker{c2, c1}, []string{n2, n1}},
	}

	for _, tc := range testCases {
		result := checkerNames(tc.have...)
		t.Run(tc.desc, func(t *testing.T) {
			if reflect.DeepEqual(tc.want, result) {
				t.Errorf("want %#v, got %#v", tc.want, result)
			}
		})
	}

}
<|endoftext|>


<commit_before>// Package wkhtmltoimage provides a simple wrapper around wkhtmltoimage (http://wkhtmltopdf.org/) binary.
package wkhtmltoimage

import (
	"bytes"
	"errors"
	"image/jpeg"
	"image/png"
	"os/exec"
	"strconv"
	"strings"
)

// ImageOptions represent the options to generate the image.
type ImageOptions struct {
	// BinaryPrefix any prefix command to wkhtmltoimage
	//
	// Any command to prefex wkhtmltoimage (e.g. xvfb-run)
	BinaryPrefix string
	// BinaryPath the path to your wkhtmltoimage binary. REQUIRED
	//
	// Must be absolute path e.g /usr/local/bin/wkhtmltoimage
	BinaryPath string
	// Input is the content to turn into an image. REQUIRED
	//
	// Can be a url (http://example.com), a local file (/tmp/example.html), or html as a string (send "-" and set the Html value)
	Input string
	// Format is the type of image to generate
	//
	// jpg, png, svg, bmp supported. Defaults to local wkhtmltoimage default
	Format string
	// Height is the height of the screen used to render in pixels.
	//
	// Default is calculated from page content. Default 0 (renders entire page top to bottom)
	Height int
	// Width is the width of the screen used to render in pixels.
	//
	// Note that this is used only as a guide line. Default 1024
	Width int
	// Quality determines the final image quality.
	//
	// Values supported between 1 and 100. Default is 94
	Quality int
	// Html is a string of html to render into and image.
	//
	// Only needed to be set if Input is set to "-"
	Html string
	// Output controls how to save or return the image.
	//
	// Leave nil to return a []byte of the image. Set to a path (/tmp/example.png) to save as a file.
	Output string
}

// GenerateImage creates an image from an input.
// It returns the image ([]byte) and any error encountered.
func GenerateImage(options *ImageOptions) ([]byte, error) {
	arr, err := buildParams(options)
	if err != nil {
		return []byte{}, err
	}

	if options.BinaryPath == "" {
		return []byte{}, errors.New("BinaryPath not set")
	}

	// cmd := new(exec.Cmd)
	// if options.BinaryPrefix != "" {
	// 	arr = append([]string{options.BinaryPath}, arr...)
	// 	cmd = exec.Command(options.BinaryPrefix, arr...)
	// } else {
	cmd := exec.Command(options.BinaryPath, arr...)
	// }

	if options.Html != "" {
		cmd.Stdin = strings.NewReader(options.Html)
	}

	output, err := cmd.CombinedOutput()

	trimmed := cleanupOutput(output, options.Format)

	return trimmed, err
}

// buildParams takes the image options set by the user and turns them into command flags for wkhtmltoimage
// It returns an array of command flags.
func buildParams(options *ImageOptions) ([]string, error) {
	a := []string{}

	if options.Input == "" {
		return []string{}, errors.New("Must provide input")
	}

	// silence extra wkhtmltoimage output
	// might want to add --javascript-delay too?
	a = append(a, "-q")
	a = append(a, "--disable-plugins")
	// a = append(a, "--use-xserver")

	a = append(a, "--format")
	if options.Format != "" {
		a = append(a, options.Format)
	} else {
		a = append(a, "png")
	}

	if options.Height != 0 {
		a = append(a, "--height")
		a = append(a, strconv.Itoa(options.Height))
	}

	if options.Width != 0 {
		a = append(a, "--width")
		a = append(a, strconv.Itoa(options.Width))
	}

	if options.Quality != 0 {
		a = append(a, "--quality")
		a = append(a, strconv.Itoa(options.Quality))
	}

	// url and output come last
	if options.Input != "-" {
		// make sure we dont pass stdin if we aren't expecting it
		options.Html = ""
	}

	a = append(a, options.Input)

	if options.Output == "" {
		a = append(a, "-")
	} else {
		a = append(a, options.Output)
	}

	return a, nil
}

func cleanupOutput(img []byte, format string) []byte {
	buf := new(bytes.Buffer)
	switch {
	case format == "png":
		decoded, err := png.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = png.Decode(bytes.NewReader(img))
		}
		png.Encode(buf, decoded)
		return buf.Bytes()
	case format == "jpg":
		decoded, err := jpeg.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = jpeg.Decode(bytes.NewReader(img))
		}
		jpeg.Encode(buf, decoded, nil)
		return buf.Bytes()
		// case format == "svg":
		// 	return img
	}
	return img
}
<commit_msg>remove rest of prefix code<commit_after>// Package wkhtmltoimage provides a simple wrapper around wkhtmltoimage (http://wkhtmltopdf.org/) binary.
package wkhtmltoimage

import (
	"bytes"
	"errors"
	"image/jpeg"
	"image/png"
	"os/exec"
	"strconv"
	"strings"
)

// ImageOptions represent the options to generate the image.
type ImageOptions struct {
	// BinaryPath the path to your wkhtmltoimage binary. REQUIRED
	//
	// Must be absolute path e.g /usr/local/bin/wkhtmltoimage
	BinaryPath string
	// Input is the content to turn into an image. REQUIRED
	//
	// Can be a url (http://example.com), a local file (/tmp/example.html), or html as a string (send "-" and set the Html value)
	Input string
	// Format is the type of image to generate
	//
	// jpg, png, svg, bmp supported. Defaults to local wkhtmltoimage default
	Format string
	// Height is the height of the screen used to render in pixels.
	//
	// Default is calculated from page content. Default 0 (renders entire page top to bottom)
	Height int
	// Width is the width of the screen used to render in pixels.
	//
	// Note that this is used only as a guide line. Default 1024
	Width int
	// Quality determines the final image quality.
	//
	// Values supported between 1 and 100. Default is 94
	Quality int
	// Html is a string of html to render into and image.
	//
	// Only needed to be set if Input is set to "-"
	Html string
	// Output controls how to save or return the image.
	//
	// Leave nil to return a []byte of the image. Set to a path (/tmp/example.png) to save as a file.
	Output string
}

// GenerateImage creates an image from an input.
// It returns the image ([]byte) and any error encountered.
func GenerateImage(options *ImageOptions) ([]byte, error) {
	arr, err := buildParams(options)
	if err != nil {
		return []byte{}, err
	}

	if options.BinaryPath == "" {
		return []byte{}, errors.New("BinaryPath not set")
	}

	cmd := exec.Command(options.BinaryPath, arr...)

	if options.Html != "" {
		cmd.Stdin = strings.NewReader(options.Html)
	}

	output, err := cmd.CombinedOutput()

	trimmed := cleanupOutput(output, options.Format)

	return trimmed, err
}

// buildParams takes the image options set by the user and turns them into command flags for wkhtmltoimage
// It returns an array of command flags.
func buildParams(options *ImageOptions) ([]string, error) {
	a := []string{}

	if options.Input == "" {
		return []string{}, errors.New("Must provide input")
	}

	// silence extra wkhtmltoimage output
	// might want to add --javascript-delay too?
	a = append(a, "-q")
	a = append(a, "--disable-plugins")

	a = append(a, "--format")
	if options.Format != "" {
		a = append(a, options.Format)
	} else {
		a = append(a, "png")
	}

	if options.Height != 0 {
		a = append(a, "--height")
		a = append(a, strconv.Itoa(options.Height))
	}

	if options.Width != 0 {
		a = append(a, "--width")
		a = append(a, strconv.Itoa(options.Width))
	}

	if options.Quality != 0 {
		a = append(a, "--quality")
		a = append(a, strconv.Itoa(options.Quality))
	}

	// url and output come last
	if options.Input != "-" {
		// make sure we dont pass stdin if we aren't expecting it
		options.Html = ""
	}

	a = append(a, options.Input)

	if options.Output == "" {
		a = append(a, "-")
	} else {
		a = append(a, options.Output)
	}

	return a, nil
}

func cleanupOutput(img []byte, format string) []byte {
	buf := new(bytes.Buffer)
	switch {
	case format == "png":
		decoded, err := png.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = png.Decode(bytes.NewReader(img))
		}
		png.Encode(buf, decoded)
		return buf.Bytes()
	case format == "jpg":
		decoded, err := jpeg.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = jpeg.Decode(bytes.NewReader(img))
		}
		jpeg.Encode(buf, decoded, nil)
		return buf.Bytes()
		// case format == "svg":
		// 	return img
	}
	return img
}
<|endoftext|>


<commit_before>package worker

import "github.com/qor/qor/resource"

type Worker struct {
	*resource.Resource
	Handler   func(job *Job)
	OnError   func(job *Job)
	OnSuccess func(job *Job)
	OnStart   func(job *Job)
	OnKill    func(job *Job)
	Adapter
}

func New(name string) Worker {
	return Worker{}
}

func (worker Worker) AllowSchedule() {
	worker.RegisterMeta(&resource.Meta{Name: "Schedule", Type: "schedule"})
}

func (worker Worker) AllMetas() []resource.Meta {
	return []resource.Meta{}
}

func (worker Worker) UseAdapter(adapter Adapter) {
	worker.Adapter = adapter
}

func (worker Worker) Listen() {
	worker.Adapter.Listen(worker)
}
<commit_msg>RunHandler for Worker<commit_after>package worker

import "github.com/qor/qor/resource"

type Worker struct {
	*resource.Resource
	Handler   func(job *Job) error
	OnError   func(job *Job)
	OnSuccess func(job *Job)
	OnStart   func(job *Job)
	OnKill    func(job *Job)
	Adapter
}

func New(name string) Worker {
	return Worker{}
}

func (worker Worker) AllowSchedule() {
	worker.RegisterMeta(&resource.Meta{Name: "Schedule", Type: "schedule"})
}

func (worker Worker) AllMetas() []resource.Meta {
	return []resource.Meta{}
}

func (worker Worker) RunHandler(job *Job) {
	defer func() {
		if r := recover(); r != nil {
			worker.OnError(job)
		}
	}()

	worker.OnStart(job)
	if err := worker.Handler(job); err == nil {
		worker.OnSuccess(job)
	} else {
		worker.OnError(job)
	}
}

func (worker Worker) UseAdapter(adapter Adapter) {
	worker.Adapter = adapter
}

func (worker Worker) Listen() {
	worker.Adapter.Listen(worker)
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", time.Now().Format(time.RFC3339), "date to start backfilling data from. Use RFC3339 format")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	formatted_date_string, err := time.Parse(time.RFC3339, start_from)
	if err != nil {
		log.Fatal("[backfill] error parsing date to start from", err)
	}
	formatted_date_string_with_tz := formatted_date_string.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<commit_msg>update backfilled to fetch based on timestamps in the database<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", "", "date to start backfilling data from. Use RFC3339 format. Default will be the tiem of the least recently updated SR in the database.")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	var fetch_from time.Time

	if start_from == "" {
		err := worker.Db.QueryRow("SELECT updated_datetime FROM service_requests ORDER BY updated_datetime ASC LIMIT 1").Scan(&fetch_from)
		if err != nil {
			log.Println("error fetching oldest SR:", err)
		}
		log.Printf("no start_from value provided, so falling back to oldest (by last update) SR in the database: %s", start_from)
	} else {
		t, err := time.Parse(time.RFC3339, start_from)
		if err != nil {
			log.Fatal("[backfill] error parsing date to start from", err)
		}
		fetch_from = t
	}

	formatted_date_string_with_tz := fetch_from.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<|endoftext|>


<commit_before>package wsutil

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"

	"github.com/gobwas/ws"
)

// DebugDialer is a wrapper around ws.Dialer. It tracks i/o of WebSocket
// handshake. That is, it gives ability to receive copied HTTP request and
// response bytes that made inside Dialer.Dial().
//
// Note that it must not be used in production applications that requires
// Dial() efficiency.
type DebugDialer struct {
	// Dialer points to a Dialer that will make Dial(). If Dialer is nil then
	// the empty Dialer will be used.
	Dialer *ws.Dialer

	// OnRequest and OnResponse are the callbacks that will be called with the
	// HTTP request and response respectively.
	OnRequest, OnResponse func([]byte)
}

// Dial connects to the url host and upgrades connection to WebSocket. It makes
// it by calling d.Dialer.Dial().
func (d *DebugDialer) Dial(ctx context.Context, urlstr string) (conn net.Conn, br *bufio.Reader, hs ws.Handshake, err error) {
	// Need to copy Dialer to prevent original object mutation.
	var dialer ws.Dialer
	if d.Dialer == nil {
		dialer = ws.Dialer{}
	} else {
		dialer = *d.Dialer
	}
	var (
		reqBuf bytes.Buffer
		resBuf bytes.Buffer

		resContentLength int64
	)
	userWrap := dialer.WrapConn
	dialer.WrapConn = func(c net.Conn) net.Conn {
		if userWrap != nil {
			c = userWrap(c)
		}

		// Save the pointer to the raw connection.
		conn = c

		var (
			r io.Reader = conn
			w io.Writer = conn
		)
		if d.OnResponse != nil {
			r = &prefetchResponseReader{
				source:        conn,
				buffer:        &resBuf,
				contentLength: &resContentLength,
			}
		}
		if d.OnRequest != nil {
			w = io.MultiWriter(conn, &reqBuf)
		}
		return rwConn{conn, r, w}
	}

	_, br, hs, err = dialer.Dial(ctx, urlstr)

	if onRequest := d.OnRequest; onRequest != nil {
		onRequest(reqBuf.Bytes())
	}
	if onResponse := d.OnResponse; onResponse != nil {
		// We must split response iniside buffered bytes from other received
		// bytes from server.
		p := resBuf.Bytes()
		n := bytes.Index(p, headEnd)
		h := n + len(headEnd)         // Head end index.
		n = h + int(resContentLength) // Body end index.

		onResponse(p[:n])

		if br != nil {
			// If br is non-nil, then it mean two things. First is that
			// handshake is OK and server has sent additional bytes – probably
			// response body or immediate sent frames. Second, the bad one, is
			// that br buffer's source is now rwConn instance from
			// above WrapConn call. It is incorrect, so we must fix it.
			var r io.Reader = conn
			if len(p) > h {
				// Buffer contains more than just HTTP headers bytes.
				r = io.MultiReader(
					bytes.NewReader(p[h:]),
					conn,
				)
			}
			br.Reset(r)
			// Must make br.Buffered() non-zero.
			br.Peek(len(p[h:]))
		}
	}

	return conn, br, hs, err
}

type rwConn struct {
	net.Conn

	r io.Reader
	w io.Writer
}

func (rwc rwConn) Read(p []byte) (int, error) {
	return rwc.r.Read(p)
}
func (rwc rwConn) Write(p []byte) (int, error) {
	return rwc.w.Write(p)
}

var headEnd = []byte("\r\n\r\n")

type prefetchResponseReader struct {
	source io.Reader
	reader io.Reader
	buffer *bytes.Buffer

	contentLength *int64
}

func (r *prefetchResponseReader) Read(p []byte) (int, error) {
	if r.reader == nil {
		resp, err := http.ReadResponse(bufio.NewReader(
			io.TeeReader(r.source, r.buffer),
		), nil)
		if err == nil {
			*r.contentLength, _ = io.Copy(ioutil.Discard, resp.Body)
			resp.Body.Close()
		}
		bts := r.buffer.Bytes()
		r.reader = io.MultiReader(
			bytes.NewReader(bts),
			r.source,
		)
	}
	return r.reader.Read(p)
}
<commit_msg>wsutil: fix comments<commit_after>package wsutil

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"

	"github.com/gobwas/ws"
)

// DebugDialer is a wrapper around ws.Dialer. It tracks i/o of WebSocket
// handshake. That is, it gives ability to receive copied HTTP request and
// response bytes that made inside Dialer.Dial().
//
// Note that it must not be used in production applications that requires
// Dial() efficiency.
type DebugDialer struct {
	// Dialer points to a Dialer that will make Dial(). If Dialer is nil then
	// the empty Dialer will be used.
	Dialer *ws.Dialer

	// OnRequest and OnResponse are the callbacks that will be called with the
	// HTTP request and response respectively.
	OnRequest, OnResponse func([]byte)
}

// Dial connects to the url host and upgrades connection to WebSocket. It makes
// it by calling d.Dialer.Dial().
func (d *DebugDialer) Dial(ctx context.Context, urlstr string) (conn net.Conn, br *bufio.Reader, hs ws.Handshake, err error) {
	// Need to copy Dialer to prevent original object mutation.
	var dialer ws.Dialer
	if d.Dialer == nil {
		dialer = ws.Dialer{}
	} else {
		dialer = *d.Dialer
	}
	var (
		reqBuf bytes.Buffer
		resBuf bytes.Buffer

		resContentLength int64
	)
	userWrap := dialer.WrapConn
	dialer.WrapConn = func(c net.Conn) net.Conn {
		if userWrap != nil {
			c = userWrap(c)
		}

		// Save the pointer to the raw connection.
		conn = c

		var (
			r io.Reader = conn
			w io.Writer = conn
		)
		if d.OnResponse != nil {
			r = &prefetchResponseReader{
				source:        conn,
				buffer:        &resBuf,
				contentLength: &resContentLength,
			}
		}
		if d.OnRequest != nil {
			w = io.MultiWriter(conn, &reqBuf)
		}
		return rwConn{conn, r, w}
	}

	_, br, hs, err = dialer.Dial(ctx, urlstr)

	if onRequest := d.OnRequest; onRequest != nil {
		onRequest(reqBuf.Bytes())
	}
	if onResponse := d.OnResponse; onResponse != nil {
		// We must split response iniside buffered bytes from other received
		// bytes from server.
		p := resBuf.Bytes()
		n := bytes.Index(p, headEnd)
		h := n + len(headEnd)         // Head end index.
		n = h + int(resContentLength) // Body end index.

		onResponse(p[:n])

		if br != nil {
			// If br is non-nil, then it mean two things. First is that
			// handshake is OK and server has sent additional bytes – probably
			// response body or immediate sent frames. Second, the bad one, is
			// that br buffer's source is now rwConn instance from
			// above WrapConn call. It is incorrect, so we must fix it.
			var r io.Reader = conn
			if len(p) > h {
				// Buffer contains more than just HTTP headers bytes.
				r = io.MultiReader(
					bytes.NewReader(p[h:]),
					conn,
				)
			}
			br.Reset(r)
			// Must make br.Buffered() to be non-zero.
			br.Peek(len(p[h:]))
		}
	}

	return conn, br, hs, err
}

type rwConn struct {
	net.Conn

	r io.Reader
	w io.Writer
}

func (rwc rwConn) Read(p []byte) (int, error) {
	return rwc.r.Read(p)
}
func (rwc rwConn) Write(p []byte) (int, error) {
	return rwc.w.Write(p)
}

var headEnd = []byte("\r\n\r\n")

type prefetchResponseReader struct {
	source io.Reader
	reader io.Reader
	buffer *bytes.Buffer

	contentLength *int64
}

func (r *prefetchResponseReader) Read(p []byte) (int, error) {
	if r.reader == nil {
		resp, err := http.ReadResponse(bufio.NewReader(
			io.TeeReader(r.source, r.buffer),
		), nil)
		if err == nil {
			*r.contentLength, _ = io.Copy(ioutil.Discard, resp.Body)
			resp.Body.Close()
		}
		bts := r.buffer.Bytes()
		r.reader = io.MultiReader(
			bytes.NewReader(bts),
			r.source,
		)
	}
	return r.reader.Read(p)
}
<|endoftext|>


<commit_before>package mettbot

import (
	a "./answers"
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	irc "github.com/fluffle/goirc/client"
	"html"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var Host *string = flag.String("host", "irc.ps0ke.de:2342", "IRC server")
var Channel *string = flag.String("channel", "#metttest", "IRC channel")
var Nick *string = flag.String("nick", "rohmett", "IRC nick")
var Longnick *string = flag.String("longnick", "Le MettBot", "IRC fullname")
var Timeformat *string = flag.String("timeformat", "2006-01-02T15:04", "Time format string (standard date: 2006-01-02T15:04:05)")
var Quotes *string = flag.String("quotes", "mett_quotes.txt", "Quote database file")
var Metts *string = flag.String("metts", "mett_metts.txt", "Metts database file")
var Offtime *int = flag.Int("offtime", 4, "Number of hours of offtopic content befor posting mett content")
var Offmessages *int = flag.Int("offmessages", 100, "Number of messages of offtopic content befor posting mett content")
var Probability *float64 = flag.Float64("probability", 0.1, "Probability that the bot ignores a command")
var Twitterregex *string = flag.String("twitterregex", "\\S*twitter\\.com\\/\\S+\\/status(es)?\\/(\\d+)\\S*", "The regex to match Twitter URLs")
var Firebird *float64 = flag.Float64("firebird", 0.05, "Probability firebird gets a question")

func init() {
	flag.Parse()
	rand.Seed(time.Now().Unix())
}

type Mettbot struct {
	*irc.Conn
	Quitted         chan bool
	QuotesPrnt      chan string
	MettsPrnt       chan string
	QuotesLinesPrnt chan int
	MettsLinesPrnt  chan int
	Input           chan string
	IsMett          chan bool
	ReallyQuit      bool
	Topics          map[string]string
	MsgSinceMett    int
}

func NewMettbot(nick string, args ...string) *Mettbot {
	bot := &Mettbot{
		irc.SimpleClient(nick, args...), // *irc.Conn
		make(chan bool),                 // Quitted
		make(chan string),               // QuotesPrnt
		make(chan string),               // MettsPrnt
		make(chan int),                  // QuotesLinesPrnt
		make(chan int),                  // MettsLinesPrnt
		make(chan string, 4),            // Input
		make(chan bool),                 // IsMett
		false,                           // ReallyQuit
		make(map[string]string), // Topics
		0,                       // MsgSinceMett
	}
	bot.EnableStateTracking()
	return bot
}

func (bot *Mettbot) Syntax(channel string) {
	bot.Notice(channel, a.RandStr(a.Syntax))
}

func (bot *Mettbot) Mentioned(channel string) {
	bot.Notice(channel, a.RandStr(a.Mention))
}

func (bot *Mettbot) Mett() {
	bot.MsgSinceMett = 0
	bot.IsMett <- true
}

func (bot *Mettbot) CheckMett() {
	for {
		select {
		case <-bot.IsMett:
		case <-time.After(time.Duration(*Offtime) * time.Hour):
			hour := time.Now().Hour()
			if hour < 1 || hour >= 8 {
				bot.Notice(*Channel, fmt.Sprintf(a.RandStr(a.MettTime), bot.GetMett(*Channel)))
			}
		}
	}
}

func (bot *Mettbot) GetMett(channel string) string {
	fi, err := os.Open(*Metts)
	if err != nil {
		log.Println(err)
		bot.Notice(channel, "Failed to open mett database")
		return "Error"
	}
	defer fi.Close()

	reader := bufio.NewReader(fi)
	lines := 0
	for {
		_, err = reader.ReadString('\n')
		if err != nil {
			break
		}
		lines++
	}

	num := rand.Intn(lines)
	mett := ""

	_, err = fi.Seek(0, 0)
	if err != nil {
		log.Println(err)
	}
	for ; num >= 0; num-- {
		mett, err = reader.ReadString('\n')
		if err == io.EOF {
			log.Println("PostMett: reached EOF")
			return "Error"
		}
		if err != nil {
			log.Println(err)
			bot.Notice(channel, "Failed to read from mett database")
			return "Error"
		}
	}
	return mett
}

func (bot *Mettbot) diffTopic(oldTopic, newTopic string) string {
	oldFile, err := ioutil.TempFile("", ".mettbotWdiffOld")
	if err != nil {
		log.Println(err)
		return ""
	}

	n, err := oldFile.WriteString(oldTopic)
	if n != len(oldTopic) || err != nil {
		log.Println(err)
		return ""
	}
	oldFile.Close()

	newFile, err := ioutil.TempFile("", ".mettbotWdiffNew")
	if err != nil {
		log.Println(err)
		return ""
	}

	n, err = newFile.WriteString(newTopic)
	if n != len(newTopic) || err != nil {
		log.Println(err)
		return ""
	}

	newFile.Close()
	defer func() {
		os.Remove(oldFile.Name())
		os.Remove(newFile.Name())
	}()

	db := "❣" // DeletionBegin
	de := "❢" // DeletionEnd
	ib := "¶" // InsertionBegin
	ie := "⁋" // InsertionEnd

	for {
		rdb := rune(rand.Intn(255-32) + 32)
		rde := rune(rand.Intn(255-32) + 32)
		rib := rune(rand.Intn(255-32) + 32)
		rie := rune(rand.Intn(255-32) + 32)

		contains := strings.ContainsRune(oldTopic, rdb)
		contains = contains || strings.ContainsRune(oldTopic, rde)
		contains = contains || strings.ContainsRune(oldTopic, rib)
		contains = contains || strings.ContainsRune(oldTopic, rie)

		contains = contains || strings.ContainsRune(newTopic, rdb)
		contains = contains || strings.ContainsRune(newTopic, rde)
		contains = contains || strings.ContainsRune(newTopic, rib)
		contains = contains || strings.ContainsRune(newTopic, rie)

		if contains == false {
			db = fmt.Sprintf("%c", rdb)
			de = fmt.Sprintf("%c", rde)
			ib = fmt.Sprintf("%c", rib)
			ie = fmt.Sprintf("%c", rie)
			break
		}
	}

	coloring := map[string]string{ // http://oreilly.com/pub/h/1953
		db: "\x035\x1F\x02",
		de: "\x0F\x0315",
		ib: "\x033\x02",
		ie: "\x0F\x0315",
	}

	cmd := exec.Command("wdiff", "-w"+db, "-x"+de, "-y"+ib, "-z"+ie, oldFile.Name(), newFile.Name())
	out, _ := cmd.Output()
	outStr := string(out)

	for n, v := range coloring {
		outStr = strings.Replace(outStr, n, v, -1)
	}

	return outStr
}

func (bot *Mettbot) GetTweet(channel, url string) {
	regex, err := regexp.Compile(*Twitterregex)
	if err != nil {
		log.Println(err)
		return
	}
	sub := regex.FindStringSubmatch(url)
	tweetUrl := fmt.Sprintf("https://api.twitter.com/1/statuses/show.json?id=%v&trim_user=false&include_entities=no", sub[2])

	resp, err := http.Get(tweetUrl)
	if err != nil {
		log.Println(err)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err)
		return
	}

	type usr struct {
		Screen_name string
	}
	type tweet struct {
		Text string
		User usr
	}

	var twt tweet
	err = json.Unmarshal(body, &twt)
	if err != nil {
		log.Println(err)
		return
	}

	tweetText := html.UnescapeString(twt.Text)

	tweetText = strings.Map( func(r rune) rune {
		if r < ' ' {
			return -1
		}
		return r
	}, tweetText)

	bot.Notice(channel, "@"+twt.User.Screen_name+": "+tweetText)
}

func (bot *Mettbot) firebird(channel string) {
	time.Sleep(time.Duration(rand.Intn(3)+3) * time.Second)
	bot.Notice(channel, a.RandStr(a.Firebird))
}

func (bot *Mettbot) WriteQuote(filename string, prnt <-chan string, linesPrnt chan<- int) {
	for message := range prnt {
		fo, err := os.OpenFile(filename, syscall.O_RDWR+syscall.O_CREAT, 0644)
		if err != nil {
			log.Println(err)
			bot.Notice(*Channel, "Couldn't open quote database")
			continue
		}
		defer fo.Close()

		foReader := bufio.NewReader(fo)
		lines := 0 //last quote has no newline

		for {
			_, err = foReader.ReadString('\n')
			if err != nil {
				break
			}
			lines++
		}
		linesPrnt <- lines

		_, err = fo.WriteString(message)
		if err != nil {
			log.Println(err)
			bot.Notice(*Channel, "Couldn't write to quote database")
		}
		fo.Close()
	}
}

func (bot *Mettbot) ReadStdin() {
	con := bufio.NewReader(os.Stdin)
	for {
		s, err := con.ReadString('\n')
		if err != nil {
			// wha?, maybe ctrl-D...
			close(bot.Input)
			break
		}
		// no point in sending empty lines down the channel
		if len(s) > 2 {
			bot.Input <- s[0 : len(s)-1]
		}
	}
}

func (bot *Mettbot) ParseStdin() {
	for cmd := range bot.Input {
		if cmd[0] == ':' {
			idx := strings.Index(cmd, " ")
			msg := cmd[idx+1:]

			switch {
			case cmd[1] == 'd':
				fmt.Printf(bot.String())
			case cmd[1] == 'f':
				if len(cmd) > 2 && cmd[2] == 'e' {
					// enable flooding
					bot.Flood = true
				} else if len(cmd) > 2 && cmd[2] == 'd' {
					// disable flooding
					bot.Flood = false
				}
				for i := 0; i < 20; i++ {
					bot.Privmsg(*Channel, "salami!1!")
				}
			case idx == -1:
				continue
			case cmd[1] == 'q':
				bot.ReallyQuit = true
				bot.Quit(msg)
			case cmd[1] == 'j':
				bot.Join(msg)
			case cmd[1] == 'p':
				bot.Part(msg)
			case cmd[1] == 'm':
				bot.Privmsg(*Channel, msg)
			case cmd[1] == 'a':
				bot.Action(*Channel, msg)
			case cmd[1] == 'n':
				bot.Notice(*Channel, msg)
			case cmd[1] == 's':
				midx := strings.Index(msg, " ")
				if midx == -1 {
					fmt.Println("Wrong Syntax")
					continue
				}
				val := msg[midx+1:]
				switch msg[:midx] {
				case "channel":
					bot.Join(val)
					*Channel = val
				case "nick":
					bot.Nick(val)
					*Nick = val
				case "quotes":
					*Quotes = val
				case "metts":
					*Metts = val
				case "offtime":
					num, err := strconv.Atoi(val)
					if err != nil {
						fmt.Println("No Number")
						continue
					}
					*Offtime = num
				case "offmessages":
					num, err := strconv.Atoi(val)
					if err != nil {
						fmt.Println("No Number")
						continue
					}
					*Offmessages = num
				case "probability":
					num, err := strconv.ParseFloat(val, 64)
					if err != nil {
						fmt.Println("No Flaot")
						continue
					}
					*Probability = num
				case "firebird":
					num, err := strconv.ParseFloat(val, 64)
					if err != nil {
						fmt.Println("No Flaot")
						continue
					}
					*Firebird = num
				default:
					fmt.Println("Unknown variable")
				}
			}
		} else {
			bot.Raw(cmd)
		}
	}
}
<commit_msg>changed newlinehandling in tweets<commit_after>package mettbot

import (
	a "./answers"
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	irc "github.com/fluffle/goirc/client"
	"html"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var Host *string = flag.String("host", "irc.ps0ke.de:2342", "IRC server")
var Channel *string = flag.String("channel", "#metttest", "IRC channel")
var Nick *string = flag.String("nick", "rohmett", "IRC nick")
var Longnick *string = flag.String("longnick", "Le MettBot", "IRC fullname")
var Timeformat *string = flag.String("timeformat", "2006-01-02T15:04", "Time format string (standard date: 2006-01-02T15:04:05)")
var Quotes *string = flag.String("quotes", "mett_quotes.txt", "Quote database file")
var Metts *string = flag.String("metts", "mett_metts.txt", "Metts database file")
var Offtime *int = flag.Int("offtime", 4, "Number of hours of offtopic content befor posting mett content")
var Offmessages *int = flag.Int("offmessages", 100, "Number of messages of offtopic content befor posting mett content")
var Probability *float64 = flag.Float64("probability", 0.1, "Probability that the bot ignores a command")
var Twitterregex *string = flag.String("twitterregex", "\\S*twitter\\.com\\/\\S+\\/status(es)?\\/(\\d+)\\S*", "The regex to match Twitter URLs")
var Firebird *float64 = flag.Float64("firebird", 0.05, "Probability firebird gets a question")

func init() {
	flag.Parse()
	rand.Seed(time.Now().Unix())
}

type Mettbot struct {
	*irc.Conn
	Quitted         chan bool
	QuotesPrnt      chan string
	MettsPrnt       chan string
	QuotesLinesPrnt chan int
	MettsLinesPrnt  chan int
	Input           chan string
	IsMett          chan bool
	ReallyQuit      bool
	Topics          map[string]string
	MsgSinceMett    int
}

func NewMettbot(nick string, args ...string) *Mettbot {
	bot := &Mettbot{
		irc.SimpleClient(nick, args...), // *irc.Conn
		make(chan bool),                 // Quitted
		make(chan string),               // QuotesPrnt
		make(chan string),               // MettsPrnt
		make(chan int),                  // QuotesLinesPrnt
		make(chan int),                  // MettsLinesPrnt
		make(chan string, 4),            // Input
		make(chan bool),                 // IsMett
		false,                           // ReallyQuit
		make(map[string]string), // Topics
		0,                       // MsgSinceMett
	}
	bot.EnableStateTracking()
	return bot
}

func (bot *Mettbot) Syntax(channel string) {
	bot.Notice(channel, a.RandStr(a.Syntax))
}

func (bot *Mettbot) Mentioned(channel string) {
	bot.Notice(channel, a.RandStr(a.Mention))
}

func (bot *Mettbot) Mett() {
	bot.MsgSinceMett = 0
	bot.IsMett <- true
}

func (bot *Mettbot) CheckMett() {
	for {
		select {
		case <-bot.IsMett:
		case <-time.After(time.Duration(*Offtime) * time.Hour):
			hour := time.Now().Hour()
			if hour < 1 || hour >= 8 {
				bot.Notice(*Channel, fmt.Sprintf(a.RandStr(a.MettTime), bot.GetMett(*Channel)))
			}
		}
	}
}

func (bot *Mettbot) GetMett(channel string) string {
	fi, err := os.Open(*Metts)
	if err != nil {
		log.Println(err)
		bot.Notice(channel, "Failed to open mett database")
		return "Error"
	}
	defer fi.Close()

	reader := bufio.NewReader(fi)
	lines := 0
	for {
		_, err = reader.ReadString('\n')
		if err != nil {
			break
		}
		lines++
	}

	num := rand.Intn(lines)
	mett := ""

	_, err = fi.Seek(0, 0)
	if err != nil {
		log.Println(err)
	}
	for ; num >= 0; num-- {
		mett, err = reader.ReadString('\n')
		if err == io.EOF {
			log.Println("PostMett: reached EOF")
			return "Error"
		}
		if err != nil {
			log.Println(err)
			bot.Notice(channel, "Failed to read from mett database")
			return "Error"
		}
	}
	return mett
}

func (bot *Mettbot) diffTopic(oldTopic, newTopic string) string {
	oldFile, err := ioutil.TempFile("", ".mettbotWdiffOld")
	if err != nil {
		log.Println(err)
		return ""
	}

	n, err := oldFile.WriteString(oldTopic)
	if n != len(oldTopic) || err != nil {
		log.Println(err)
		return ""
	}
	oldFile.Close()

	newFile, err := ioutil.TempFile("", ".mettbotWdiffNew")
	if err != nil {
		log.Println(err)
		return ""
	}

	n, err = newFile.WriteString(newTopic)
	if n != len(newTopic) || err != nil {
		log.Println(err)
		return ""
	}

	newFile.Close()
	defer func() {
		os.Remove(oldFile.Name())
		os.Remove(newFile.Name())
	}()

	db := "❣" // DeletionBegin
	de := "❢" // DeletionEnd
	ib := "¶" // InsertionBegin
	ie := "⁋" // InsertionEnd

	for {
		rdb := rune(rand.Intn(255-32) + 32)
		rde := rune(rand.Intn(255-32) + 32)
		rib := rune(rand.Intn(255-32) + 32)
		rie := rune(rand.Intn(255-32) + 32)

		contains := strings.ContainsRune(oldTopic, rdb)
		contains = contains || strings.ContainsRune(oldTopic, rde)
		contains = contains || strings.ContainsRune(oldTopic, rib)
		contains = contains || strings.ContainsRune(oldTopic, rie)

		contains = contains || strings.ContainsRune(newTopic, rdb)
		contains = contains || strings.ContainsRune(newTopic, rde)
		contains = contains || strings.ContainsRune(newTopic, rib)
		contains = contains || strings.ContainsRune(newTopic, rie)

		if contains == false {
			db = fmt.Sprintf("%c", rdb)
			de = fmt.Sprintf("%c", rde)
			ib = fmt.Sprintf("%c", rib)
			ie = fmt.Sprintf("%c", rie)
			break
		}
	}

	coloring := map[string]string{ // http://oreilly.com/pub/h/1953
		db: "\x035\x1F\x02",
		de: "\x0F\x0315",
		ib: "\x033\x02",
		ie: "\x0F\x0315",
	}

	cmd := exec.Command("wdiff", "-w"+db, "-x"+de, "-y"+ib, "-z"+ie, oldFile.Name(), newFile.Name())
	out, _ := cmd.Output()
	outStr := string(out)

	for n, v := range coloring {
		outStr = strings.Replace(outStr, n, v, -1)
	}

	return outStr
}

func (bot *Mettbot) GetTweet(channel, url string) {
	regex, err := regexp.Compile(*Twitterregex)
	if err != nil {
		log.Println(err)
		return
	}
	sub := regex.FindStringSubmatch(url)
	tweetUrl := fmt.Sprintf("https://api.twitter.com/1/statuses/show.json?id=%v&trim_user=false&include_entities=no", sub[2])

	resp, err := http.Get(tweetUrl)
	if err != nil {
		log.Println(err)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err)
		return
	}

	type usr struct {
		Screen_name string
	}
	type tweet struct {
		Text string
		User usr
	}

	var twt tweet
	err = json.Unmarshal(body, &twt)
	if err != nil {
		log.Println(err)
		return
	}

	tweetText := html.UnescapeString(twt.Text)

	tweetText = strings.Map(func(r rune) rune {
		if r < ' ' {
			return ' '
		}
		return r
	}, tweetText)

	bot.Notice(channel, "@"+twt.User.Screen_name+": "+tweetText)
}

func (bot *Mettbot) firebird(channel string) {
	time.Sleep(time.Duration(rand.Intn(3)+3) * time.Second)
	bot.Notice(channel, a.RandStr(a.Firebird))
}

func (bot *Mettbot) WriteQuote(filename string, prnt <-chan string, linesPrnt chan<- int) {
	for message := range prnt {
		fo, err := os.OpenFile(filename, syscall.O_RDWR+syscall.O_CREAT, 0644)
		if err != nil {
			log.Println(err)
			bot.Notice(*Channel, "Couldn't open quote database")
			continue
		}
		defer fo.Close()

		foReader := bufio.NewReader(fo)
		lines := 0 //last quote has no newline

		for {
			_, err = foReader.ReadString('\n')
			if err != nil {
				break
			}
			lines++
		}
		linesPrnt <- lines

		_, err = fo.WriteString(message)
		if err != nil {
			log.Println(err)
			bot.Notice(*Channel, "Couldn't write to quote database")
		}
		fo.Close()
	}
}

func (bot *Mettbot) ReadStdin() {
	con := bufio.NewReader(os.Stdin)
	for {
		s, err := con.ReadString('\n')
		if err != nil {
			// wha?, maybe ctrl-D...
			close(bot.Input)
			break
		}
		// no point in sending empty lines down the channel
		if len(s) > 2 {
			bot.Input <- s[0 : len(s)-1]
		}
	}
}

func (bot *Mettbot) ParseStdin() {
	for cmd := range bot.Input {
		if cmd[0] == ':' {
			idx := strings.Index(cmd, " ")
			msg := cmd[idx+1:]

			switch {
			case cmd[1] == 'd':
				fmt.Printf(bot.String())
			case cmd[1] == 'f':
				if len(cmd) > 2 && cmd[2] == 'e' {
					// enable flooding
					bot.Flood = true
				} else if len(cmd) > 2 && cmd[2] == 'd' {
					// disable flooding
					bot.Flood = false
				}
				for i := 0; i < 20; i++ {
					bot.Privmsg(*Channel, "salami!1!")
				}
			case idx == -1:
				continue
			case cmd[1] == 'q':
				bot.ReallyQuit = true
				bot.Quit(msg)
			case cmd[1] == 'j':
				bot.Join(msg)
			case cmd[1] == 'p':
				bot.Part(msg)
			case cmd[1] == 'm':
				bot.Privmsg(*Channel, msg)
			case cmd[1] == 'a':
				bot.Action(*Channel, msg)
			case cmd[1] == 'n':
				bot.Notice(*Channel, msg)
			case cmd[1] == 's':
				midx := strings.Index(msg, " ")
				if midx == -1 {
					fmt.Println("Wrong Syntax")
					continue
				}
				val := msg[midx+1:]
				switch msg[:midx] {
				case "channel":
					bot.Join(val)
					*Channel = val
				case "nick":
					bot.Nick(val)
					*Nick = val
				case "quotes":
					*Quotes = val
				case "metts":
					*Metts = val
				case "offtime":
					num, err := strconv.Atoi(val)
					if err != nil {
						fmt.Println("No Number")
						continue
					}
					*Offtime = num
				case "offmessages":
					num, err := strconv.Atoi(val)
					if err != nil {
						fmt.Println("No Number")
						continue
					}
					*Offmessages = num
				case "probability":
					num, err := strconv.ParseFloat(val, 64)
					if err != nil {
						fmt.Println("No Flaot")
						continue
					}
					*Probability = num
				case "firebird":
					num, err := strconv.ParseFloat(val, 64)
					if err != nil {
						fmt.Println("No Flaot")
						continue
					}
					*Firebird = num
				default:
					fmt.Println("Unknown variable")
				}
			}
		} else {
			bot.Raw(cmd)
		}
	}
}
<|endoftext|>