<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/alecthomas/template"

	"gopkg.in/errgo.v1"
	"gopkg.in/yaml.v2"
)

var userFlag = flag.String("user", "", "default github user")
var projectFlag = flag.String("project", "", "default github project")
var editFlag = flag.Bool("edit", false, "edit config")
var setupFlag = flag.Bool("setup", false, "run project setup")

type session struct {
	Name        string            `yaml:"name"`
	SetupScript string            `yaml:"setup-script"`
	Environment map[string]string `yaml:"environment"`
	Cwd         string            `yaml:"cwd"`
	Windows     []window          `yaml:"windows"`
	Focus       string            `yaml:"focus"`
}

type window struct {
	Name       string   `yaml:"name"`
	Command    string   `yaml:"command"`
	Cwd        string   `yaml:"cwd"`
	Keystrokes []string `yaml:"keystrokes"`
}

func die(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, errgo.Details(err))
		os.Exit(1)
	}
}

func main() {
	die(run())
}

func run() error {
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Fprintln(os.Stderr, "usage: %s <session.yaml file>", os.Args[0])
		return errgo.New("missing session file argument")
	}

	name, err := locateSession(flag.Arg(0))
	if os.IsNotExist(err) || *editFlag {
		*setupFlag = true
		err = newSessionFile(flag.Arg(0))
	}
	if err != nil {
		return errgo.Mask(err)
	}
	if *editFlag {
		os.Exit(0)
	}

	session, err := newSession(name)
	if err != nil {
		return errgo.Mask(err)
	}

	for k, v := range session.Environment {
		os.Setenv(k, os.ExpandEnv(v))
	}
	if _, err := os.Stat(session.Cwd); os.IsNotExist(err) {
		*setupFlag = true
	}

	log.Println(os.Environ())

	if *setupFlag {
		err = session.setupScript()
		if err != nil {
			return errgo.Notef(err, "failed to execute setup script")
		}
	}

	err = session.create()
	if err != nil {
		return errgo.Mask(err)
	}

	var focus int
	for i, window := range session.Windows {
		if i > 0 {
			err = session.createWindow(i, &window)
			if err != nil {
				return errgo.Mask(err)
			}
		}

		err = session.sendKeys(&window)
		if err != nil {
			return errgo.Mask(err)
		}

		if session.Focus == window.Name {
			focus = i
		}
	}
	err = session.focus(focus)
	if err != nil {
		return errgo.Mask(err)
	}

	err = session.tmux("attach", "-t", session.Name)
	return errgo.Mask(err)
}

func locateSession(name string) (string, error) {
	if info, err := os.Stat(name); err == nil && !info.IsDir() {
		return name, nil
	} else if err != nil && !os.IsNotExist(err) {
		return "", errgo.Notef(err, "failed to open session file %q", name)
	}

	configDir := os.Getenv("XDG_CONFIG_HOME")
	if configDir == "" {
		configDir = filepath.Join(os.Getenv("HOME"), ".config")
	}

	tmuxgConfigDir := filepath.Join(configDir, "tmuxg")
	err := os.MkdirAll(tmuxgConfigDir, 0644)
	if err != nil {
		return "", errgo.Notef(err, "failed to create config directory %q", tmuxgConfigDir)
	}

	confPath := filepath.Join(tmuxgConfigDir, name+".yaml")
	if _, err := os.Stat(confPath); err != nil {
		if os.IsNotExist(err) {
			return "", err
		}
		return "", errgo.Notef(err, "failed to resolve session %q file %q", name, confPath)
	}
	return confPath, nil
}

const newTemplateContents = `
# Name of the session. Probably don't mess with this.
name: {{.Name}}

# Environment variables set for the tmux session.
environment:
  GOPATH: ${HOME}/go/{{.Name}}

# Current working directory for the tmux session. May use environment variables
# declared above.
cwd: ${GOPATH}/src/github.com/{{.User}}/{{.Project}}

# Script to run the first time this session starts, or when specifically
# invoked with -setup. Try to make this script idempotent.
setup-script: |
    #!/bin/bash
    mkdir -p ${GOPATH}
    go get -d github.com/{{.User}}/{{.Project}}/...

# Windows to create in the tmux session and what to run in each.
windows:
  - name: editor
    command: vim
    keystrokes:
      - \n
  - name: shell
focus: editor
`

var newTemplate = template.Must(template.New("new-conf").Parse(newTemplateContents))

func newSessionFile(name string) error {
	configDir := os.Getenv("XDG_CONFIG_HOME")
	if configDir == "" {
		configDir = filepath.Join(os.Getenv("HOME"), ".config")
	}

	tmuxgConfigDir := filepath.Join(configDir, "tmuxg")
	err := os.MkdirAll(tmuxgConfigDir, 0644)
	if err != nil {
		return errgo.Notef(err, "failed to create config directory %q", tmuxgConfigDir)
	}
	confPath := filepath.Join(tmuxgConfigDir, name+".yaml")
	err = createSessionFile(name, confPath)
	if err != nil {
		defer os.Remove(confPath)
	}
	return err
}

func createSessionFile(name, confPath string) error {
	if _, err := os.Stat(confPath); os.IsNotExist(err) {
		err = func() error {
			f, err := os.Create(confPath)
			if err != nil {
				return errgo.Notef(err, "failed to open config file %q for writing", confPath)
			}
			defer f.Close()

			user := *userFlag
			if user == "" {
				user = os.Getenv("USER")
			}
			project := *projectFlag
			if project == "" {
				project = name
			}

			err = newTemplate.Execute(f, struct {
				Name, User, Project string
			}{
				Name:    name,
				User:    user,
				Project: project,
			})
			if err != nil {
				return errgo.Notef(err, "failed to write config file %q", confPath)
			}
			return nil
		}()
		if err != nil {
			return err
		}
	}
	cmd := exec.Command("vim", confPath)
	cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
	err := cmd.Run()
	if err != nil {
		return errgo.Notef(err, "editor exited with error")
	}
	_, err = newSession(confPath)
	return err
}

func newSession(confPath string) (*session, error) {
	var s session

	contents, err := ioutil.ReadFile(confPath)
	if err != nil {
		return nil, errgo.Notef(err, "failed to read session file")
	}
	err = yaml.Unmarshal(contents, &s)
	if err != nil {
		return nil, errgo.Notef(err, "failed to parse session file")
	}

	for i := range s.Windows {
		if s.Windows[i].Command == "" {
			s.Windows[i].Command = "bash"
		}
	}
	return &s, nil
}

func (s *session) tmux(args ...string) error {
	c := exec.Command("tmux", append([]string{"-L", s.Name}, args...)...)
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	c.Dir = os.ExpandEnv(s.Cwd)
	log.Printf("%v", c)
	return errgo.Mask(c.Run())
}

func (s *session) setupScript() error {
	if s.SetupScript == "" {
		return nil
	}

	f, err := ioutil.TempFile("", "tmuxg-setup")
	if err != nil {
		return errgo.Notef(err, "failed to create temporary file for script")
	}
	defer os.Remove(f.Name())
	defer f.Close()

	_, err = fmt.Fprintf(f, "%s", strings.TrimSpace(s.SetupScript))
	if err != nil {
		return errgo.Notef(err, "failed to write temporary script file")
	}
	err = f.Close()
	if err != nil {
		return errgo.Notef(err, "error closing temporary script file")
	}
	err = os.Chmod(f.Name(), 0700)
	if err != nil {
		return errgo.Notef(err, "error setting temporary script executable")
	}

	c := exec.Command(f.Name())
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	return errgo.Mask(c.Run())
}

func (s *session) create() error {
	if len(s.Windows) == 0 {
		return errgo.New("no windows configured for this session!")
	}
	err := s.tmux("new-session", "-d", "-s", s.Name,
		os.ExpandEnv(s.Windows[0].Command))
	if err != nil {
		return errgo.Notef(err, "failed to start tmux session")
	}

	var keys []string
	for k, v := range s.Environment {
		err = s.tmux("set-environment", "-t", s.Name, k, os.ExpandEnv(v))
		if err != nil {
			return errgo.Notef(err, "warning: failed to set environment variable %q", k)
		}
		keys = append(keys, k)
	}

	return nil
}

func (s *session) createWindow(i int, w *window) error {
	cwd := w.Cwd
	if cwd == "" {
		cwd = s.Cwd
	}
	cwd = os.ExpandEnv(cwd)
	err := s.tmux("new-window", "-d", "-t", fmt.Sprintf("%s:%d", s.Name, i),
		"-c", cwd, os.ExpandEnv(w.Command))
	if err != nil {
		return errgo.Notef(err, "failed to create window %q", w.Name)
	}
	return nil
}

func (s *session) sendKeys(w *window) error {
	if len(w.Keystrokes) > 0 {
		err := s.tmux(append([]string{"send-keys"}, w.Keystrokes...)...)
		if err != nil {
			return errgo.Notef(err, "failed to send keystrokes to window %q", w.Name)
		}
	}
	return nil
}

func (s *session) focus(i int) error {
	err := s.tmux("select-window", "-t", fmt.Sprintf("%s:%d", s.Name, i))
	if err != nil {
		return errgo.Notef(err, "failed to set window focus")
	}
	return nil
}
<commit_msg>Just use stdlib text/template.<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"

	"gopkg.in/errgo.v1"
	"gopkg.in/yaml.v2"
)

var userFlag = flag.String("user", "", "default github user")
var projectFlag = flag.String("project", "", "default github project")
var editFlag = flag.Bool("edit", false, "edit config")
var setupFlag = flag.Bool("setup", false, "run project setup")

type session struct {
	Name        string            `yaml:"name"`
	SetupScript string            `yaml:"setup-script"`
	Environment map[string]string `yaml:"environment"`
	Cwd         string            `yaml:"cwd"`
	Windows     []window          `yaml:"windows"`
	Focus       string            `yaml:"focus"`
}

type window struct {
	Name       string   `yaml:"name"`
	Command    string   `yaml:"command"`
	Cwd        string   `yaml:"cwd"`
	Keystrokes []string `yaml:"keystrokes"`
}

func die(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, errgo.Details(err))
		os.Exit(1)
	}
}

func main() {
	die(run())
}

func run() error {
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Fprintln(os.Stderr, "usage: %s <session.yaml file>", os.Args[0])
		return errgo.New("missing session file argument")
	}

	name, err := locateSession(flag.Arg(0))
	if os.IsNotExist(err) || *editFlag {
		*setupFlag = true
		err = newSessionFile(flag.Arg(0))
	}
	if err != nil {
		return errgo.Mask(err)
	}
	if *editFlag {
		os.Exit(0)
	}

	session, err := newSession(name)
	if err != nil {
		return errgo.Mask(err)
	}

	for k, v := range session.Environment {
		os.Setenv(k, os.ExpandEnv(v))
	}
	if _, err := os.Stat(session.Cwd); os.IsNotExist(err) {
		*setupFlag = true
	}

	log.Println(os.Environ())

	if *setupFlag {
		err = session.setupScript()
		if err != nil {
			return errgo.Notef(err, "failed to execute setup script")
		}
	}

	err = session.create()
	if err != nil {
		return errgo.Mask(err)
	}

	var focus int
	for i, window := range session.Windows {
		if i > 0 {
			err = session.createWindow(i, &window)
			if err != nil {
				return errgo.Mask(err)
			}
		}

		err = session.sendKeys(&window)
		if err != nil {
			return errgo.Mask(err)
		}

		if session.Focus == window.Name {
			focus = i
		}
	}
	err = session.focus(focus)
	if err != nil {
		return errgo.Mask(err)
	}

	err = session.tmux("attach", "-t", session.Name)
	return errgo.Mask(err)
}

func locateSession(name string) (string, error) {
	if info, err := os.Stat(name); err == nil && !info.IsDir() {
		return name, nil
	} else if err != nil && !os.IsNotExist(err) {
		return "", errgo.Notef(err, "failed to open session file %q", name)
	}

	configDir := os.Getenv("XDG_CONFIG_HOME")
	if configDir == "" {
		configDir = filepath.Join(os.Getenv("HOME"), ".config")
	}

	tmuxgConfigDir := filepath.Join(configDir, "tmuxg")
	err := os.MkdirAll(tmuxgConfigDir, 0644)
	if err != nil {
		return "", errgo.Notef(err, "failed to create config directory %q", tmuxgConfigDir)
	}

	confPath := filepath.Join(tmuxgConfigDir, name+".yaml")
	if _, err := os.Stat(confPath); err != nil {
		if os.IsNotExist(err) {
			return "", err
		}
		return "", errgo.Notef(err, "failed to resolve session %q file %q", name, confPath)
	}
	return confPath, nil
}

const newTemplateContents = `
# Name of the session. Probably don't mess with this.
name: {{.Name}}

# Environment variables set for the tmux session.
environment:
  GOPATH: ${HOME}/go/{{.Name}}

# Current working directory for the tmux session. May use environment variables
# declared above.
cwd: ${GOPATH}/src/github.com/{{.User}}/{{.Project}}

# Script to run the first time this session starts, or when specifically
# invoked with -setup. Try to make this script idempotent.
setup-script: |
    #!/bin/bash
    mkdir -p ${GOPATH}
    go get -d github.com/{{.User}}/{{.Project}}/...

# Windows to create in the tmux session and what to run in each.
windows:
  - name: editor
    command: vim
    keystrokes:
      - \n
  - name: shell
focus: editor
`

var newTemplate = template.Must(template.New("new-conf").Parse(newTemplateContents))

func newSessionFile(name string) error {
	configDir := os.Getenv("XDG_CONFIG_HOME")
	if configDir == "" {
		configDir = filepath.Join(os.Getenv("HOME"), ".config")
	}

	tmuxgConfigDir := filepath.Join(configDir, "tmuxg")
	err := os.MkdirAll(tmuxgConfigDir, 0644)
	if err != nil {
		return errgo.Notef(err, "failed to create config directory %q", tmuxgConfigDir)
	}
	confPath := filepath.Join(tmuxgConfigDir, name+".yaml")
	err = createSessionFile(name, confPath)
	if err != nil {
		defer os.Remove(confPath)
	}
	return err
}

func createSessionFile(name, confPath string) error {
	if _, err := os.Stat(confPath); os.IsNotExist(err) {
		err = func() error {
			f, err := os.Create(confPath)
			if err != nil {
				return errgo.Notef(err, "failed to open config file %q for writing", confPath)
			}
			defer f.Close()

			user := *userFlag
			if user == "" {
				user = os.Getenv("USER")
			}
			project := *projectFlag
			if project == "" {
				project = name
			}

			err = newTemplate.Execute(f, struct {
				Name, User, Project string
			}{
				Name:    name,
				User:    user,
				Project: project,
			})
			if err != nil {
				return errgo.Notef(err, "failed to write config file %q", confPath)
			}
			return nil
		}()
		if err != nil {
			return err
		}
	}
	cmd := exec.Command("vim", confPath)
	cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
	err := cmd.Run()
	if err != nil {
		return errgo.Notef(err, "editor exited with error")
	}
	_, err = newSession(confPath)
	return err
}

func newSession(confPath string) (*session, error) {
	var s session

	contents, err := ioutil.ReadFile(confPath)
	if err != nil {
		return nil, errgo.Notef(err, "failed to read session file")
	}
	err = yaml.Unmarshal(contents, &s)
	if err != nil {
		return nil, errgo.Notef(err, "failed to parse session file")
	}

	for i := range s.Windows {
		if s.Windows[i].Command == "" {
			s.Windows[i].Command = "bash"
		}
	}
	return &s, nil
}

func (s *session) tmux(args ...string) error {
	c := exec.Command("tmux", append([]string{"-L", s.Name}, args...)...)
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	c.Dir = os.ExpandEnv(s.Cwd)
	log.Printf("%v", c)
	return errgo.Mask(c.Run())
}

func (s *session) setupScript() error {
	if s.SetupScript == "" {
		return nil
	}

	f, err := ioutil.TempFile("", "tmuxg-setup")
	if err != nil {
		return errgo.Notef(err, "failed to create temporary file for script")
	}
	defer os.Remove(f.Name())
	defer f.Close()

	_, err = fmt.Fprintf(f, "%s", strings.TrimSpace(s.SetupScript))
	if err != nil {
		return errgo.Notef(err, "failed to write temporary script file")
	}
	err = f.Close()
	if err != nil {
		return errgo.Notef(err, "error closing temporary script file")
	}
	err = os.Chmod(f.Name(), 0700)
	if err != nil {
		return errgo.Notef(err, "error setting temporary script executable")
	}

	c := exec.Command(f.Name())
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	return errgo.Mask(c.Run())
}

func (s *session) create() error {
	if len(s.Windows) == 0 {
		return errgo.New("no windows configured for this session!")
	}
	err := s.tmux("new-session", "-d", "-s", s.Name,
		os.ExpandEnv(s.Windows[0].Command))
	if err != nil {
		return errgo.Notef(err, "failed to start tmux session")
	}

	var keys []string
	for k, v := range s.Environment {
		err = s.tmux("set-environment", "-t", s.Name, k, os.ExpandEnv(v))
		if err != nil {
			return errgo.Notef(err, "warning: failed to set environment variable %q", k)
		}
		keys = append(keys, k)
	}

	return nil
}

func (s *session) createWindow(i int, w *window) error {
	cwd := w.Cwd
	if cwd == "" {
		cwd = s.Cwd
	}
	cwd = os.ExpandEnv(cwd)
	err := s.tmux("new-window", "-d", "-t", fmt.Sprintf("%s:%d", s.Name, i),
		"-c", cwd, os.ExpandEnv(w.Command))
	if err != nil {
		return errgo.Notef(err, "failed to create window %q", w.Name)
	}
	return nil
}

func (s *session) sendKeys(w *window) error {
	if len(w.Keystrokes) > 0 {
		err := s.tmux(append([]string{"send-keys"}, w.Keystrokes...)...)
		if err != nil {
			return errgo.Notef(err, "failed to send keystrokes to window %q", w.Name)
		}
	}
	return nil
}

func (s *session) focus(i int) error {
	err := s.tmux("select-window", "-t", fmt.Sprintf("%s:%d", s.Name, i))
	if err != nil {
		return errgo.Notef(err, "failed to set window focus")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"sphere2cubeGo/cache"
	"time"
	"flag"
	"sphere2cubeGo/worker"
	"sphere2cubeGo/saver"
)

var (
	tileNames = []string{
		worker.TileUp,
		worker.TileDown,
		worker.TileFront,
		worker.TileRight,
		worker.TileBack,
		worker.TileLeft,
	}
	tileSize          = 1024
	originalImagePath = ""
	outPutDir         = "./build"
)

func main() {

	tileSizeCmd := flag.Int("s", tileSize, "Size in px of final tile")
	originalImagePathCmd := flag.String("i", "", "Path to input equirectangular panorama")
	outPutDirCmd := flag.String("o", outPutDir, "Path to output directory")

	flag.Parse()

	tileSize = *tileSizeCmd
	originalImagePath = *originalImagePathCmd
	outPutDir = *outPutDirCmd

	if originalImagePath == "" {
		flag.PrintDefaults()
		os.Exit(2)
	}

	_, err := os.Stat(originalImagePath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Fatalf("%v not found", originalImagePath)
		}
	}

	done := make(chan worker.TileResult)
	timeStart := time.Now()
	cacheResult := cache.CacheAnglesHandler(tileSize)
	for _, tileName := range tileNames {
		tile := worker.Tile{TileName: tileName, TileSize: tileSize}
		go worker.Worker(tile, cacheResult, originalImagePath, done)
	}

	for range tileNames {
		tileResult := <-done
		err = saver.SaveTile(tileResult, outPutDir)

		if err != nil {
			log.Fatal(err.Error())
			os.Exit(2)
		}

		log.Printf("Process for tile %v --> finished", tileResult.Tile.TileName)
	}

	timeFinish := time.Now()
	duration := timeFinish.Sub(timeStart)
	log.Printf("Time to render: %v seconds", duration.Seconds())
}
<commit_msg>go fmt<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"sphere2cubeGo/cache"
	"sphere2cubeGo/saver"
	"sphere2cubeGo/worker"
	"time"
)

var (
	tileNames = []string{
		worker.TileUp,
		worker.TileDown,
		worker.TileFront,
		worker.TileRight,
		worker.TileBack,
		worker.TileLeft,
	}
	tileSize          = 1024
	originalImagePath = ""
	outPutDir         = "./build"
)

func main() {

	tileSizeCmd := flag.Int("s", tileSize, "Size in px of final tile")
	originalImagePathCmd := flag.String("i", "", "Path to input equirectangular panorama")
	outPutDirCmd := flag.String("o", outPutDir, "Path to output directory")

	flag.Parse()

	tileSize = *tileSizeCmd
	originalImagePath = *originalImagePathCmd
	outPutDir = *outPutDirCmd

	if originalImagePath == "" {
		flag.PrintDefaults()
		os.Exit(2)
	}

	_, err := os.Stat(originalImagePath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Fatalf("%v not found", originalImagePath)
		}
	}

	done := make(chan worker.TileResult)
	timeStart := time.Now()
	cacheResult := cache.CacheAnglesHandler(tileSize)
	for _, tileName := range tileNames {
		tile := worker.Tile{TileName: tileName, TileSize: tileSize}
		go worker.Worker(tile, cacheResult, originalImagePath, done)
	}

	for range tileNames {
		tileResult := <-done
		err = saver.SaveTile(tileResult, outPutDir)

		if err != nil {
			log.Fatal(err.Error())
			os.Exit(2)
		}

		log.Printf("Process for tile %v --> finished", tileResult.Tile.TileName)
	}

	timeFinish := time.Now()
	duration := timeFinish.Sub(timeStart)
	log.Printf("Time to render: %v seconds", duration.Seconds())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/golang-lru"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

const (
	LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT = time.Minute
	LETSENCRYPT_PRODUCTION_API_URL         = "https://acme-v01.api.letsencrypt.org/directory"
	LETSENCRYPT_STAGING_API_URL            = "https://acme-staging.api.letsencrypt.org/directory"
	PRIVATE_KEY_BITS                       = 2048
	TRY_COUNT                              = 10
	RETRY_SLEEP                            = time.Second * 5
	STATE_FILEMODE                         = 0600
)

var (
	bindTo                 = flag.String("bind-to", ":443", "")
	targetPort             = flag.Int("target-port", 80, "")
	targetConnTimeout      = flag.Duration("target-conn-timeout", time.Second, "")
	acmeApiUrl             = flag.String("acme-server", LETSENCRYPT_PRODUCTION_API_URL, "")
	acmeTestServer         = flag.Bool("test", false, "Use test lets encrypt server instead of <acme-server>")
	certDir                = flag.String("cert-dir", "certificates", `Directory for save cached certificates. Set cert-dir=- for disable save certs`)
	certMemCount           = flag.Int("in-memory-cnt", 10000, "How many count of certs cache in memory for prevent parse it from file")
	stateFilePath          = flag.String("state-file", "state.json", "Path to save some state data, for example account key")
	proxyMode              = flag.String("proxy-mode", "http", "Proxy-mode after tls handle (http|tcp).")
	httpRealIPHeader       = flag.String("real-ip-header", "X-Real-IP", "The header will contain original IP of remote connection. It can be few headers, separated by comma.")
	additionalHeadersParam = flag.String("additional-headers", "X-Forwarded-Proto=https", "Additional headers for proxied requests. Several headers separated by comma.")
	logLevel               = flag.String("loglevel", "warning", "fatal|error|warning|info|debug")
	logOutput              = flag.String("logout", "-", "Path to logout. Special: '-' (without quotes) - stderr")
	versionPrint           = flag.Bool("version", false, "print version and exit.")
)

var (
	VERSION           = "unversioned" // need be var becouse it redefine by --ldflags "-X main.VERSION" during autobuild
	realIPHeaderNames [][]byte        // IP headers, generated by the proxy, included real IP address
	cutHeaders        [][]byte        // internal - all headers, that cut from request (upper case).
	additionalHeaders []byte          // prepared additional headers
	acmeService       *acmeStruct
)

type stateStruct struct {
	PrivateKey *rsa.PrivateKey
	changed    bool
}

func main() {
	var err error

	flag.Parse()

	if *versionPrint {
		fmt.Println(strings.TrimSpace(VERSION))
		return
	}

	// Init
	if *logOutput != "-" {
		fLog, err := os.OpenFile(*logOutput, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
		if fLog != nil {
			defer fLog.Close()
		}
		if err == nil {
			logrus.SetOutput(fLog)
		} else {
			logrus.Errorf("Can't open logfile '%v': %v", *logOutput, err)
		}
	}
	logrus.SetLevel(logrus.WarnLevel)
	switch *logLevel {
	case "fatal":
		logrus.SetLevel(logrus.FatalLevel)
	case "error":
		logrus.SetLevel(logrus.ErrorLevel)
	case "warning":
		logrus.SetLevel(logrus.WarnLevel)
	case "info":
		logrus.SetLevel(logrus.InfoLevel)
	case "debug":
		logrus.SetLevel(logrus.DebugLevel)
	default:
		logrus.Errorf("Use default loglevel '%v', becouse unknow level: '%v'", logrus.GetLevel(), *logLevel)
	}
	logrus.Infof("Use log level: '%v'", logrus.GetLevel())

	if *proxyMode != "http" && *proxyMode != "tcp" {
		logrus.Panicf("Unknow proxy mode: %v", *proxyMode)
	}
	logrus.Infof("Proxy mode: %v", *proxyMode)

	for _, line := range strings.Split(*httpRealIPHeader, ",") {
		line = strings.TrimSpace(line)
		if line != "" {
			realIPHeaderNames = append(realIPHeaderNames, []byte(line))
			cutHeaders = append(realIPHeaderNames, []byte(strings.ToUpper(line)))
		}
	}

	for _, addHeader := range strings.Split(*additionalHeadersParam, ",") {
		headerParts := strings.SplitN(addHeader, "=", 2)
		if len(headerParts) > 0 {
			cutHeaders = append(cutHeaders, []byte(strings.ToUpper(headerParts[0])))
		}
		buf := &bytes.Buffer{}
		buf.WriteString(headerParts[0])
		buf.WriteByte(':')
		if len(headerParts) == 2 {
			buf.WriteString(headerParts[1])
		}
		buf.WriteString("\r\n")
		additionalHeaders = append(additionalHeaders, buf.Bytes()...)
	}

	localIPs = getLocalIPs()
	acmeService = &acmeStruct{}
	if *certMemCount > 0 {
		logrus.Infof("Create memory cache for '%v' certificates", *certMemCount)
		certMemCache, err = lru.New(*certMemCount)
		if err != nil {
			logrus.Errorf("Can't create memory cache:", err)
			certMemCache = nil
		}
	} else {
		logrus.Info("Memory cache turned off")
	}
	if *certDir == "-" {
		*certDir = ""
	}

	// init service
	var state stateStruct
	stateBytes, err := ioutil.ReadFile(*stateFilePath)
	if err == nil {
		err = json.Unmarshal(stateBytes, &state)
		if err != nil {
			logrus.Errorf("Can't parse state file '%v': %v", *stateFilePath, err)
		}
	} else {
		logrus.Errorf("Can't read state file '%v': %v", *stateFilePath, err)
	}

	acmeService = &acmeStruct{}
	if *acmeTestServer {
		acmeService.serverAddress = LETSENCRYPT_STAGING_API_URL
	} else {
		acmeService.serverAddress = *acmeApiUrl
	}

	if state.PrivateKey == nil {
		logrus.Info("Generate private keys")
		state.PrivateKey, err = rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
		state.changed = true
		if err != nil {
			logrus.Panic("Can't generate private key")
		}
	} else {
		logrus.Debugf("Skip generate keys - it was read from state")
	}

	saveState(state)

	acmeService.privateKey = state.PrivateKey

	acmeService.Init()
	acmeService.RegisterEnsure(context.TODO())

	// Start listen
	tcpAddr, err := net.ResolveTCPAddr("tcp", *bindTo)
	if err != nil {
		logrus.Panicf("Can't resolve bind-to address '%v': %v", *bindTo, err)
	}
	logrus.Errorf("Start listen: %v", tcpAddr)

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(err)
	}
	for {
		tcpConn, err := listener.AcceptTCP()
		if err != nil {
			panic(err)
		}
		go handleTcpConnection(tcpConn)
	}
}

func certificateGet(clientHello *tls.ClientHelloInfo) (cert *tls.Certificate, err error) {
	domain := strings.ToLower(clientHello.ServerName)
	if logrus.GetLevel() >= logrus.DebugLevel {
		logrus.Debugf("Required certificate for domain '%v'", domain)
	}

	if strings.HasSuffix(domain, ACME_DOMAIN_SUFFIX) {
		// force generate new certificate, without caching.
		return acmeService.CreateCertificate(domain)
	}

	cert = certificateCacheGet(domain)

	switch {
	case cert != nil && cert.Leaf != nil && cert.Leaf.NotAfter.Before(time.Now()):
		logrus.Warnf("Expired certificate got from cache for domain '%v'", domain)
		// pass to create new certificate.
	case cert != nil:
		return cert, nil
	default:
		// pass
	}

	cert, err = acmeService.CreateCertificate(domain)
	if err == nil {
		certificateCachePut(domain, cert)
	}
	return cert, err
}

func handleTcpConnection(in *net.TCPConn) {
	target, err := getTargetConn(in)
	if err != nil {
		logrus.Errorf("Can't get target IP/port for '%v': %v", target.String(), err)
		return
	}

	// handle ssl
	tlsConfig := tls.Config{
		GetCertificate: certificateGet,
		MinVersion:     tls.VersionSSL30,
	}
	tlsConn := tls.Server(in, &tlsConfig)
	err = tlsConn.Handshake()
	logrus.Debug("tls ciper:", tlsConn.ConnectionState().CipherSuite)
	if err == nil {
		logrus.Debug("Handshake for incoming:", tlsConn.RemoteAddr().String())
	} else {
		logrus.Infof("Error in tls handshake from '%v':%v", tlsConn.RemoteAddr(), err)
	}

	startProxy(target, tlsConn)
}

func getLocalIPs() (res []net.IP) {
	bindAddr, _ := net.ResolveTCPAddr("tcp", *bindTo)
	if bindAddr.IP.IsUnspecified() || len(bindAddr.IP) == 0 {
		addresses, err := net.InterfaceAddrs()
		if err != nil {
			log.Panic("Can't get local ip addresses:", err)
		}
		res = make([]net.IP, 0, len(addresses))
		for _, addr := range addresses {
			logrus.Info("Local ip:", addr.String())
			ip, _, err := net.ParseCIDR(addr.String())
			if err == nil {
				res = append(res, ip)
			} else {
				logrus.Errorf("Can't parse local ip '%v': %v", addr.String(), err)
			}
		}
	} else {
		res = []net.IP{bindAddr.IP}
	}
	if logrus.GetLevel() >= logrus.InfoLevel {
		ipStrings := make([]string, len(res))
		for i, addr := range res {
			ipStrings[i] = addr.String()
		}
		logrus.Info("Local ip:", ipStrings)
	}
	return res
}

func getTargetConn(in *net.TCPConn) (net.TCPAddr, error) {
	targetAddrP, err := net.ResolveTCPAddr("tcp", in.LocalAddr().String())
	if err != nil {
		logrus.Errorf("Can't resolve local addr '%v': %v", in.LocalAddr().String(), err)
		return net.TCPAddr{}, err
	}
	targetAddrP.Port = *targetPort
	return *targetAddrP, nil
}

func saveState(state stateStruct) {
	if state.changed {
		logrus.Infof("Saving state to '%v'", *stateFilePath)
	} else {
		logrus.Debug("Skip save state becouse it isn't changed")
		return
	}
	stateBytes, err := json.MarshalIndent(&state, "", "    ")
	if err != nil {
		logrus.Errorf("Can't save state to file '%v': %v", *stateFilePath, err)
		return
	}
	err = ioutil.WriteFile(*stateFilePath+".new", stateBytes, STATE_FILEMODE)
	if err != nil {
		logrus.Errorf("Error while write state bytes to file '%v': %v", *stateFilePath+".new", err)
		return
	}

	if _, err := os.Stat(*stateFilePath); !os.IsNotExist(err) {
		logrus.Infof("Rename current state file '%v' -> '%v'", *stateFilePath, *stateFilePath+".old")

		err = os.Rename(*stateFilePath, *stateFilePath+".old")
		if err != nil {
			logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath, *stateFilePath+".old", err)
		}
	} else {
		logrus.Infof("Create new state file '%v'", *stateFilePath)
	}

	err = os.Rename(*stateFilePath+".new", *stateFilePath)
	if err != nil {
		logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath+".new", *stateFilePath, err)
	}
}<commit_msg>+ --non-cert-domains - ignore domains by regexp Close https://git.1gb.ru/rekby/lets-proxy/issues/8<commit_after>package main

import (
	"bytes"
	"context"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/golang-lru"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"
	"regexp"
	"errors"
)

const (
	LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT = time.Minute
	LETSENCRYPT_PRODUCTION_API_URL         = "https://acme-v01.api.letsencrypt.org/directory"
	LETSENCRYPT_STAGING_API_URL            = "https://acme-staging.api.letsencrypt.org/directory"
	PRIVATE_KEY_BITS                       = 2048
	TRY_COUNT                              = 10
	RETRY_SLEEP                            = time.Second * 5
	STATE_FILEMODE                         = 0600
)

var (
	bindTo                 = flag.String("bind-to", ":443", "")
	targetPort             = flag.Int("target-port", 80, "")
	targetConnTimeout      = flag.Duration("target-conn-timeout", time.Second, "")
	acmeApiUrl             = flag.String("acme-server", LETSENCRYPT_PRODUCTION_API_URL, "")
	acmeTestServer         = flag.Bool("test", false, "Use test lets encrypt server instead of <acme-server>")
	certDir                = flag.String("cert-dir", "certificates", `Directory for save cached certificates. Set cert-dir=- for disable save certs`)
	certMemCount           = flag.Int("in-memory-cnt", 10000, "How many count of certs cache in memory for prevent parse it from file")
	stateFilePath          = flag.String("state-file", "state.json", "Path to save some state data, for example account key")
	proxyMode              = flag.String("proxy-mode", "http", "Proxy-mode after tls handle (http|tcp).")
	httpRealIPHeader       = flag.String("real-ip-header", "X-Real-IP", "The header will contain original IP of remote connection. It can be few headers, separated by comma.")
	additionalHeadersParam = flag.String("additional-headers", "X-Forwarded-Proto=https", "Additional headers for proxied requests. Several headers separated by comma.")
	logLevel               = flag.String("loglevel", "warning", "fatal|error|warning|info|debug")
	logOutput              = flag.String("logout", "-", "Path to logout. Special: '-' (without quotes) - stderr")
	versionPrint           = flag.Bool("version", false, "print version and exit.")
	nonCertDomains         = flag.String("non-cert-domains", "", "No obtain certificate for mathed domains. Regexpes separated by comma.")
)

var (
	VERSION           = "unversioned" // need be var becouse it redefine by --ldflags "-X main.VERSION" during autobuild
	realIPHeaderNames [][]byte        // IP headers, generated by the proxy, included real IP address
	cutHeaders        [][]byte        // internal - all headers, that cut from request (upper case).
	additionalHeaders []byte          // prepared additional headers
	acmeService       *acmeStruct
	nonCertDomainsRegexps []*regexp.Regexp
)

type stateStruct struct {
	PrivateKey *rsa.PrivateKey
	changed    bool
}

func main() {
	var err error

	flag.Parse()

	if *versionPrint {
		fmt.Println(strings.TrimSpace(VERSION))
		return
	}

	// Init
	if *logOutput != "-" {
		fLog, err := os.OpenFile(*logOutput, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
		if fLog != nil {
			defer fLog.Close()
		}
		if err == nil {
			logrus.SetOutput(fLog)
		} else {
			logrus.Errorf("Can't open logfile '%v': %v", *logOutput, err)
		}
	}
	logrus.SetLevel(logrus.WarnLevel)
	switch *logLevel {
	case "fatal":
		logrus.SetLevel(logrus.FatalLevel)
	case "error":
		logrus.SetLevel(logrus.ErrorLevel)
	case "warning":
		logrus.SetLevel(logrus.WarnLevel)
	case "info":
		logrus.SetLevel(logrus.InfoLevel)
	case "debug":
		logrus.SetLevel(logrus.DebugLevel)
	default:
		logrus.Errorf("Use default loglevel '%v', becouse unknow level: '%v'", logrus.GetLevel(), *logLevel)
	}
	logrus.Infof("Use log level: '%v'", logrus.GetLevel())

	if *proxyMode != "http" && *proxyMode != "tcp" {
		logrus.Panicf("Unknow proxy mode: %v", *proxyMode)
	}
	logrus.Infof("Proxy mode: %v", *proxyMode)

	for _, ignoreDomain := range strings.Split(*nonCertDomains, ",") {
		ignoreDomain = strings.TrimSpace(ignoreDomain)
		if ignoreDomain == "" {
			continue
		}
		ignoreDomainRE, err := regexp.Compile(ignoreDomain)
		if err != nil {
			logrus.Errorf("Bad ignore domain regexp '%v': %v", ignoreDomain, err)
		}
		if ignoreDomainRE != nil {
			nonCertDomainsRegexps = append(nonCertDomainsRegexps, ignoreDomainRE)
		}
	}
	if logrus.GetLevel() >= logrus.InfoLevel {
		regexps := []string{}
		for _, re := range nonCertDomainsRegexps {
			regexps = append(regexps, re.String())
		}
		logrus.Info("Non cert domain regexps: ", "['" + strings.Join(regexps, "', '") + "']")
	}

	for _, line := range strings.Split(*httpRealIPHeader, ",") {
		line = strings.TrimSpace(line)
		if line != "" {
			realIPHeaderNames = append(realIPHeaderNames, []byte(line))
			cutHeaders = append(realIPHeaderNames, []byte(strings.ToUpper(line)))
		}
	}

	for _, addHeader := range strings.Split(*additionalHeadersParam, ",") {
		headerParts := strings.SplitN(addHeader, "=", 2)
		if len(headerParts) > 0 {
			cutHeaders = append(cutHeaders, []byte(strings.ToUpper(headerParts[0])))
		}
		buf := &bytes.Buffer{}
		buf.WriteString(headerParts[0])
		buf.WriteByte(':')
		if len(headerParts) == 2 {
			buf.WriteString(headerParts[1])
		}
		buf.WriteString("\r\n")
		additionalHeaders = append(additionalHeaders, buf.Bytes()...)
	}

	localIPs = getLocalIPs()
	acmeService = &acmeStruct{}
	if *certMemCount > 0 {
		logrus.Infof("Create memory cache for '%v' certificates", *certMemCount)
		certMemCache, err = lru.New(*certMemCount)
		if err != nil {
			logrus.Errorf("Can't create memory cache:", err)
			certMemCache = nil
		}
	} else {
		logrus.Info("Memory cache turned off")
	}
	if *certDir == "-" {
		*certDir = ""
	}

	// init service
	var state stateStruct
	stateBytes, err := ioutil.ReadFile(*stateFilePath)
	if err == nil {
		err = json.Unmarshal(stateBytes, &state)
		if err != nil {
			logrus.Errorf("Can't parse state file '%v': %v", *stateFilePath, err)
		}
	} else {
		logrus.Errorf("Can't read state file '%v': %v", *stateFilePath, err)
	}

	acmeService = &acmeStruct{}
	if *acmeTestServer {
		acmeService.serverAddress = LETSENCRYPT_STAGING_API_URL
	} else {
		acmeService.serverAddress = *acmeApiUrl
	}

	if state.PrivateKey == nil {
		logrus.Info("Generate private keys")
		state.PrivateKey, err = rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
		state.changed = true
		if err != nil {
			logrus.Panic("Can't generate private key")
		}
	} else {
		logrus.Debugf("Skip generate keys - it was read from state")
	}

	saveState(state)

	acmeService.privateKey = state.PrivateKey

	acmeService.Init()
	acmeService.RegisterEnsure(context.TODO())

	// Start listen
	tcpAddr, err := net.ResolveTCPAddr("tcp", *bindTo)
	if err != nil {
		logrus.Panicf("Can't resolve bind-to address '%v': %v", *bindTo, err)
	}
	logrus.Errorf("Start listen: %v", tcpAddr)

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(err)
	}
	for {
		tcpConn, err := listener.AcceptTCP()
		if err != nil {
			panic(err)
		}
		go handleTcpConnection(tcpConn)
	}
}

func certificateGet(clientHello *tls.ClientHelloInfo) (cert *tls.Certificate, err error) {
	domain := strings.ToLower(clientHello.ServerName)
	if logrus.GetLevel() >= logrus.DebugLevel {
		logrus.Debugf("Required certificate for domain '%v'", domain)
	}

	if strings.HasSuffix(domain, ACME_DOMAIN_SUFFIX) {
		// force generate new certificate, without caching.
		return acmeService.CreateCertificate(domain)
	}

	cert = certificateCacheGet(domain)

	switch {
	case cert != nil && cert.Leaf != nil && cert.Leaf.NotAfter.Before(time.Now()):
		logrus.Warnf("Expired certificate got from cache for domain '%v'", domain)
		// pass to create new certificate.
	case cert != nil:
		return cert, nil
	default:
		// pass
	}

	for _, re := range nonCertDomainsRegexps {
		if re.MatchString(domain){
			logrus.Debugf("Ignore obtain cert for domain '%v' by regexp '%v'", domain, re.String())
			return nil, errors.New("Ignore cert obtain by regexp")
		}
	}

	cert, err = acmeService.CreateCertificate(domain)
	if err == nil {
		certificateCachePut(domain, cert)
	}
	return cert, err
}

func handleTcpConnection(in *net.TCPConn) {
	target, err := getTargetConn(in)
	if err != nil {
		logrus.Errorf("Can't get target IP/port for '%v': %v", target.String(), err)
		return
	}

	// handle ssl
	tlsConfig := tls.Config{
		GetCertificate: certificateGet,
		MinVersion:     tls.VersionSSL30,
	}
	tlsConn := tls.Server(in, &tlsConfig)
	err = tlsConn.Handshake()
	logrus.Debug("tls ciper:", tlsConn.ConnectionState().CipherSuite)
	if err == nil {
		logrus.Debug("Handshake for incoming:", tlsConn.RemoteAddr().String())
	} else {
		logrus.Infof("Error in tls handshake from '%v':%v", tlsConn.RemoteAddr(), err)
	}

	startProxy(target, tlsConn)
}

func getLocalIPs() (res []net.IP) {
	bindAddr, _ := net.ResolveTCPAddr("tcp", *bindTo)
	if bindAddr.IP.IsUnspecified() || len(bindAddr.IP) == 0 {
		addresses, err := net.InterfaceAddrs()
		if err != nil {
			logrus.Panic("Can't get local ip addresses:", err)
		}
		res = make([]net.IP, 0, len(addresses))
		for _, addr := range addresses {
			logrus.Info("Local ip:", addr.String())
			ip, _, err := net.ParseCIDR(addr.String())
			if err == nil {
				res = append(res, ip)
			} else {
				logrus.Errorf("Can't parse local ip '%v': %v", addr.String(), err)
			}
		}
	} else {
		res = []net.IP{bindAddr.IP}
	}
	if logrus.GetLevel() >= logrus.InfoLevel {
		ipStrings := make([]string, len(res))
		for i, addr := range res {
			ipStrings[i] = addr.String()
		}
		logrus.Info("Local ip:", ipStrings)
	}
	return res
}

func getTargetConn(in *net.TCPConn) (net.TCPAddr, error) {
	targetAddrP, err := net.ResolveTCPAddr("tcp", in.LocalAddr().String())
	if err != nil {
		logrus.Errorf("Can't resolve local addr '%v': %v", in.LocalAddr().String(), err)
		return net.TCPAddr{}, err
	}
	targetAddrP.Port = *targetPort
	return *targetAddrP, nil
}

func saveState(state stateStruct) {
	if state.changed {
		logrus.Infof("Saving state to '%v'", *stateFilePath)
	} else {
		logrus.Debug("Skip save state becouse it isn't changed")
		return
	}
	stateBytes, err := json.MarshalIndent(&state, "", "    ")
	if err != nil {
		logrus.Errorf("Can't save state to file '%v': %v", *stateFilePath, err)
		return
	}
	err = ioutil.WriteFile(*stateFilePath+".new", stateBytes, STATE_FILEMODE)
	if err != nil {
		logrus.Errorf("Error while write state bytes to file '%v': %v", *stateFilePath+".new", err)
		return
	}

	if _, err := os.Stat(*stateFilePath); !os.IsNotExist(err) {
		logrus.Infof("Rename current state file '%v' -> '%v'", *stateFilePath, *stateFilePath+".old")

		err = os.Rename(*stateFilePath, *stateFilePath+".old")
		if err != nil {
			logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath, *stateFilePath+".old", err)
		}
	} else {
		logrus.Infof("Create new state file '%v'", *stateFilePath)
	}

	err = os.Rename(*stateFilePath+".new", *stateFilePath)
	if err != nil {
		logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath+".new", *stateFilePath, err)
	}
}
<|endoftext|>


<commit_before>package robimq

import (
	"github.com/streadway/amqp"
	"log"
)

func Connect(amqpURI string) (ch *amqp.Connection) {
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		log.Printf("ERROR: Could not connect to RabbitMQ server - %q", err)
		return nil
	} else {
		log.Printf("Connected to RabbitMQ server - %s", amqpURI)
		return connection
	}
}

func Consume(conn *amqp.Connection, exchange, queue, consumerTag string) (chan string, error) {

	data := make(chan string)

	ch, err := conn.Channel()
	if err != nil {
		log.Printf("ERROR: Could not declare channel %q", err)
		return nil, err
	} else {

		log.Printf("Got Channel")
		defer ch.Close()

		err := ch.ExchangeDeclare(
			exchange, // name
			"topic",  // type
			true,     // durable
			false,    // auto-deleted
			false,    // internal
			false,    // noWait
			nil,      // arguments
		)
		if err != nil {
			log.Printf("ERROR: Could not declare Exchange [%s] %q", exchange, err)
			return nil, err
		} else {

			q, err := ch.QueueDeclare(
				queue, // name
				true,  // durable
				false, // delete when unused
				false, // exclusive
				false, // no-wait
				nil,   // arguments
			)
			if err != nil {
				log.Printf("ERROR: Could not declare queue [%s] %q", queue, err)
				return nil, err
			} else {

				err = ch.QueueBind(
					q.Name,   // queue name
					"image",  // routing key @TODO
					exchange, // exchange
					false,
					nil,
				)
				if err != nil {
					log.Printf("ERROR: Could not bind [%s] queue to [%s] exhange %q", queue, exchange, err)
					return nil, err
				} else {

					log.Printf("Queue %s declared", queue)
					err = ch.Qos(
						1,     // prefetch count
						0,     // prefetch size
						false, // global
					)
					if err != nil {
						log.Printf("ERROR: %q", err)
						return nil, err
					}
					msgs, err := ch.Consume(
						q.Name,      // queue
						consumerTag, // consumer
						false,       // auto-ack
						false,       // exclusive
						false,       // no-local
						false,       // no-wait
						nil,         // args
					)
					if err != nil {
						log.Printf("ERROR: Could not consume messages on [%s] queue %q", queue, err)
						return nil, err
					}

					forever := make(chan bool)

					go func() {
						for d := range msgs {
							d.Ack(false)
							//downloadFile(string(d.Body))
							data <- string(d.Body)
						}
					}()
					log.Printf("Consumer tag %s", consumerTag)
					<-forever
					return data, nil
				}

			}

		}

	}
	return data, nil
}
<commit_msg>FML, gotta run quick stash - not working<commit_after>package robimq

import (
	"github.com/streadway/amqp"
	"log"
)

func Connect(amqpURI string) (ch *amqp.Connection) {
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		log.Printf("ERROR: Could not connect to RabbitMQ server - %q", err)
		return nil
	} else {
		log.Printf("Connected to RabbitMQ server - %s", amqpURI)
		return connection
	}
}

func declareExchange(ch *amqp.Channel, exchange string) error {
	err := ch.ExchangeDeclare(
		exchange, // name
		"topic",  // type
		true,     // durable
		false,    // auto-deleted
		false,    // internal
		false,    // noWait
		nil,      // arguments
	)
	return err
}

func declareQueue(ch *amqp.Channel, queue string) (q amqp.Queue, err error) {
	q, err := ch.QueueDeclare(
		queue, // name
		true,  // durable
		false, // delete when unused
		false, // exclusive
		false, // no-wait
		nil,   // arguments
	)
	return q, err
}

func Consume(conn *amqp.Connection, exchange, queue, consumerTag string) (chan string, error) {

	data := make(chan string)

	ch, err := conn.Channel()
	if err != nil {
		log.Printf("ERROR: Could not declare channel %q", err)
		return nil, err
	} else {

		log.Printf("Got Channel")
		defer ch.Close()

		err := declareExchange(ch, exchange)
		if err != nil {
			log.Printf("ERROR: Could not declare Exchange [%s] %q", exchange, err)
			return nil, err
		} else {

			q, err := declareQueue(ch, queue)
			if err != nil {
				log.Printf("ERROR: Could not declare queue [%s] %q", queue, err)
				return nil, err
			} else {

				err = ch.QueueBind(
					q.Name,   // queue name
					"image",  // routing key @TODO
					exchange, // exchange
					false,
					nil,
				)
				if err != nil {
					log.Printf("ERROR: Could not bind [%s] queue to [%s] exhange %q", queue, exchange, err)
					return nil, err
				} else {

					log.Printf("Queue %s declared", queue)
					err = ch.Qos(
						1,     // prefetch count
						0,     // prefetch size
						false, // global
					)
					if err != nil {
						log.Printf("ERROR: %q", err)
						return nil, err
					}
					msgs, err := ch.Consume(
						q.Name,      // queue
						consumerTag, // consumer
						false,       // auto-ack
						false,       // exclusive
						false,       // no-local
						false,       // no-wait
						nil,         // args
					)
					if err != nil {
						log.Printf("ERROR: Could not consume messages on [%s] queue %q", queue, err)
						return nil, err
					}

					forever := make(chan bool)

					go func() {
						for d := range msgs {
							d.Ack(false)
							//downloadFile(string(d.Body))
							data <- string(d.Body)
						}
					}()
					log.Printf("Consumer tag %s", consumerTag)
					<-forever
					return data, nil
				}

			}

		}

	}
	return data, nil
}
<|endoftext|>


<commit_before>//go:generate esc -o static.go -prefix static static

package main

import (
	"encoding/json"
	"flag"
	"github.com/fatih/structs"
	"github.com/nfnt/resize"
	"gopkg.in/yaml.v2"
	"image/jpeg"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"time"
)

var photos = make(map[string]struct{})
var dirToMetadata = make(map[string]metadata)
var photoDir string
var cacheDir string

type metadata struct {
	Event        string
	Photographer string
	Date         string
	Location     string
	Directory    string `json:"-" structs:"directory"`
}

func handleFilters(w http.ResponseWriter, r *http.Request) {
	keys := []string{"Event", "Photographer", "Date", "Location"}

	filters := make(map[string]map[string]struct{})
	for _, k := range keys {
		filters[k] = make(map[string]struct{})
	}

	for _, meta := range dirToMetadata {
		m := structs.Map(meta)
		for _, k := range keys {
			if _, ok := m[k]; ok {
				filters[k][m[k].(string)] = struct{}{}
			}
		}
	}

	jFilters := make(map[string][]string)
	for _, k := range keys {
		jFilters[k] = make([]string, 0)
	}

	for i := range filters {
		for j := range filters[i] {
			jFilters[i] = append(jFilters[i], j)
		}
		sort.Strings(jFilters[i])
		if len(jFilters[i]) == 0 {
			delete(jFilters, i)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	if err := enc.Encode(jFilters); err != nil {
		log.Print(err)
		return
	}
}

func handleFilter(w http.ResponseWriter, r *http.Request) {
	filterToEligible := make(map[string][]string)

	filters := make(map[string][]string)
	for filterKey := range r.URL.Query() {
		filters[filterKey] = r.URL.Query()[filterKey]
	}

	// find all directories with matching filters
	for filterKey, filterValues := range filters {
		for dir, meta := range dirToMetadata {
			m := structs.Map(meta)
			for k, v := range m {
				if filterKey == k {
					for _, filterValue := range filterValues {
						if filterValue == v {
							filterToEligible[filterKey] = append(filterToEligible[filterKey], dir)
						}
					}
				}
			}
		}
	}

	// find intersection of all filter matches
	eligible := make(map[string]bool)

	for _, dirs := range filterToEligible {
		for _, dir := range dirs {
			count := 0
			for _, dirTests := range filterToEligible {
				for _, dirTest := range dirTests {
					if dir == dirTest {
						count += 1
					}
				}
			}
			if count == len(filterToEligible) {
				eligible[dir] = true
			}
		}
	}

	eligiblePhotos := make([]string, 0)
	for photo := range photos {
		dir := filepath.Dir(photo)
		if _, ok := eligible[dir]; ok {
			eligiblePhotos = append(eligiblePhotos, photo)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	if err := enc.Encode(eligiblePhotos); err != nil {
		log.Print(err)
		return
	}
}

func handlePhoto(w http.ResponseWriter, r *http.Request) {
	if filepath.Ext(r.URL.Path) != ".jpg" {
		http.Error(w, http.StatusText(404), 404)
		return
	}

	// serve photo
	photoFilepath := filepath.Join(photoDir, r.URL.Path[len("/photo/"):])
	http.ServeFile(w, r, photoFilepath)
}

func handleThumb(w http.ResponseWriter, r *http.Request) {
	if filepath.Ext(r.URL.Path) != ".jpg" {
		http.Error(w, http.StatusText(404), 404)
		return
	}

	photoFilepath := filepath.Join(photoDir, r.URL.Path[len("/cache/"):])
	cacheFilepath := filepath.Join(cacheDir, r.URL.Path[len("/cache/"):])

	if _, err := os.Stat(cacheFilepath); err != nil {
		if os.IsNotExist(err) {
			// create thumbnail

			file, err := os.Open(photoFilepath)
			if err != nil {
				log.Print(err)
				return
			}

			img, err := jpeg.Decode(file)
			if err != nil {
				log.Print(err)
				return
			}
			file.Close()

			m := resize.Thumbnail(500, 1000, img, resize.Bicubic)

			parentDir, _ := filepath.Split(cacheFilepath)
			err = os.MkdirAll(parentDir, os.FileMode(0764))
			if err != nil {
				log.Print(err)
				return
			}

			out, err := os.Create(cacheFilepath)
			if err != nil {
				log.Print(err)
				return
			}
			defer out.Close()

			jpeg.Encode(out, m, nil)
		} else {
			log.Print(err)
			return
		}
	}

	// serve thumbnail
	http.ServeFile(w, r, cacheFilepath)
}

func walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}

	dir, filename := filepath.Split(path)
	dir, err = filepath.Rel(photoDir, dir)
	if err != nil {
		return err
	}

	if filename == "metadata.yaml" {
		m := metadata{Directory: dir}

		data, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}

		err = yaml.Unmarshal(data, &m)
		dirToMetadata[dir] = m
	} else if filepath.Ext(filename) == ".jpg" {
		filename = filepath.Join(dir, filename)
		photos[filename] = struct{}{}
	}
	return nil
}

func walk() {
	for {
		filepath.Walk(photoDir, walkFunc)
		time.Sleep(5 * time.Second)
	}
}

func main() {
	flag.StringVar(&photoDir, "photoDir", "photos", "path to photos")
	flag.StringVar(&cacheDir, "cacheDir", "cache", "path to photo cache directory")
	flag.Parse()

	go walk()

	http.Handle("/", http.FileServer(FS(false))) // esc for static content. true uses local files, false uses embedded
	http.HandleFunc("/photo/", handlePhoto)
	http.HandleFunc("/thumb/", handleThumb)

	http.HandleFunc("/api/filters", handleFilters)
	http.HandleFunc("/api/filter", handleFilter)
	http.ListenAndServe("127.0.0.1:8000", nil)
}
<commit_msg>Sort photos so they are returned in a consistent order<commit_after>//go:generate esc -o static.go -prefix static static

package main

import (
	"encoding/json"
	"flag"
	"github.com/fatih/structs"
	"github.com/nfnt/resize"
	"gopkg.in/yaml.v2"
	"image/jpeg"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"time"
)

var photos = make(map[string]struct{})
var dirToMetadata = make(map[string]metadata)
var photoDir string
var cacheDir string

type metadata struct {
	Event        string
	Photographer string
	Date         string
	Location     string
	Directory    string `json:"-" structs:"directory"`
}

func handleFilters(w http.ResponseWriter, r *http.Request) {
	keys := []string{"Event", "Photographer", "Date", "Location"}

	filters := make(map[string]map[string]struct{})
	for _, k := range keys {
		filters[k] = make(map[string]struct{})
	}

	for _, meta := range dirToMetadata {
		m := structs.Map(meta)
		for _, k := range keys {
			if _, ok := m[k]; ok {
				filters[k][m[k].(string)] = struct{}{}
			}
		}
	}

	jFilters := make(map[string][]string)
	for _, k := range keys {
		jFilters[k] = make([]string, 0)
	}

	for i := range filters {
		for j := range filters[i] {
			jFilters[i] = append(jFilters[i], j)
		}
		sort.Strings(jFilters[i])
		if len(jFilters[i]) == 0 {
			delete(jFilters, i)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	if err := enc.Encode(jFilters); err != nil {
		log.Print(err)
		return
	}
}

func handleFilter(w http.ResponseWriter, r *http.Request) {
	filterToEligible := make(map[string][]string)

	filters := make(map[string][]string)
	for filterKey := range r.URL.Query() {
		filters[filterKey] = r.URL.Query()[filterKey]
	}

	// find all directories with matching filters
	for filterKey, filterValues := range filters {
		for dir, meta := range dirToMetadata {
			m := structs.Map(meta)
			for k, v := range m {
				if filterKey == k {
					for _, filterValue := range filterValues {
						if filterValue == v {
							filterToEligible[filterKey] = append(filterToEligible[filterKey], dir)
						}
					}
				}
			}
		}
	}

	// find intersection of all filter matches
	eligible := make(map[string]bool)

	for _, dirs := range filterToEligible {
		for _, dir := range dirs {
			count := 0
			for _, dirTests := range filterToEligible {
				for _, dirTest := range dirTests {
					if dir == dirTest {
						count += 1
					}
				}
			}
			if count == len(filterToEligible) {
				eligible[dir] = true
			}
		}
	}

	eligiblePhotos := make([]string, 0)
	for photo := range photos {
		dir := filepath.Dir(photo)
		if _, ok := eligible[dir]; ok {
			eligiblePhotos = append(eligiblePhotos, photo)
		}
	}
	sort.Strings(eligiblePhotos)

	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	if err := enc.Encode(eligiblePhotos); err != nil {
		log.Print(err)
		return
	}
}

func handlePhoto(w http.ResponseWriter, r *http.Request) {
	if filepath.Ext(r.URL.Path) != ".jpg" {
		http.Error(w, http.StatusText(404), 404)
		return
	}

	// serve photo
	photoFilepath := filepath.Join(photoDir, r.URL.Path[len("/photo/"):])
	http.ServeFile(w, r, photoFilepath)
}

func handleThumb(w http.ResponseWriter, r *http.Request) {
	if filepath.Ext(r.URL.Path) != ".jpg" {
		http.Error(w, http.StatusText(404), 404)
		return
	}

	photoFilepath := filepath.Join(photoDir, r.URL.Path[len("/cache/"):])
	cacheFilepath := filepath.Join(cacheDir, r.URL.Path[len("/cache/"):])

	if _, err := os.Stat(cacheFilepath); err != nil {
		if os.IsNotExist(err) {
			// create thumbnail

			file, err := os.Open(photoFilepath)
			if err != nil {
				log.Print(err)
				return
			}

			img, err := jpeg.Decode(file)
			if err != nil {
				log.Print(err)
				return
			}
			file.Close()

			m := resize.Thumbnail(500, 1000, img, resize.Bicubic)

			parentDir, _ := filepath.Split(cacheFilepath)
			err = os.MkdirAll(parentDir, os.FileMode(0764))
			if err != nil {
				log.Print(err)
				return
			}

			out, err := os.Create(cacheFilepath)
			if err != nil {
				log.Print(err)
				return
			}
			defer out.Close()

			jpeg.Encode(out, m, nil)
		} else {
			log.Print(err)
			return
		}
	}

	// serve thumbnail
	http.ServeFile(w, r, cacheFilepath)
}

func walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}

	dir, filename := filepath.Split(path)
	dir, err = filepath.Rel(photoDir, dir)
	if err != nil {
		return err
	}

	if filename == "metadata.yaml" {
		m := metadata{Directory: dir}

		data, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}

		err = yaml.Unmarshal(data, &m)
		dirToMetadata[dir] = m
	} else if filepath.Ext(filename) == ".jpg" {
		filename = filepath.Join(dir, filename)
		photos[filename] = struct{}{}
	}
	return nil
}

func walk() {
	for {
		filepath.Walk(photoDir, walkFunc)
		time.Sleep(5 * time.Second)
	}
}

func main() {
	flag.StringVar(&photoDir, "photoDir", "photos", "path to photos")
	flag.StringVar(&cacheDir, "cacheDir", "cache", "path to photo cache directory")
	flag.Parse()

	go walk()

	http.Handle("/", http.FileServer(FS(false))) // esc for static content. true uses local files, false uses embedded
	http.HandleFunc("/photo/", handlePhoto)
	http.HandleFunc("/thumb/", handleThumb)

	http.HandleFunc("/api/filters", handleFilters)
	http.HandleFunc("/api/filter", handleFilter)
	http.ListenAndServe("127.0.0.1:8000", nil)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"

	"github.com/mvdan/gibot/irc"
	"github.com/mvdan/gibot/site/gitlab"
)

var (
	configPath = flag.String("c", "gibot.json", "path to json config file")
)

var config struct {
	Nick  string   `json:"nick"`
	Chans []string `json:"chans"`
	Repos []struct {
		Name    string   `json:"name"`
		Url     string   `json:"url"`
		Aliases []string `json:"aliases"`
	} `json:"repos"`
}

func main() {
	flag.Parse()
	if err := loadConfig(); err != nil {
		log.Fatalf("Could not load config: %v", err)
	}
	log.Printf("nick  = %s", config.Nick)
	log.Printf("chans = %s", strings.Join(config.Chans, ", "))

	var repos []gitlab.Repo
	knownAliases := make(map[string]struct{})
	for _, r := range config.Repos {
		aliases := append(r.Aliases, r.Name)
		repos = append(repos, gitlab.NewRepo(r.Name,
			r.Url, aliases...))
		for _, a := range aliases {
			if _, e := knownAliases[a]; e {
				log.Fatalf("Alias '%s' is not unique!", a)
			}
			knownAliases[a] = struct{}{}
		}
	}

	l := &listener{
		repos: repos,
		allRe: joinRegexes(repos),
	}

	log.Printf("Connecting...")
	c, err := irc.Connect(config.Nick, config.Chans)
	if err != nil {
		log.Fatalf("Could not connect to IRC: %v", err)
	}

	l.client = c
	go l.listen()

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	<-sigc
	log.Printf("Quitting.")
	c.Quit()
}

func loadConfig() error {
	configFile, err := os.Open(*configPath)
	if err != nil {
		return err
	}
	return json.NewDecoder(configFile).Decode(&config)
}

func joinRegexes(repos []gitlab.Repo) *regexp.Regexp {
	var all []string
	for _, r := range repos {
		all = append(all, r.IssuesRe.String())
		all = append(all, r.PullsRe.String())
	}
	allRe := regexp.MustCompile(`(` + strings.Join(all, `|`) + `)`)
	allRe.Longest()
	return allRe
}

type listener struct {
	repos  []gitlab.Repo
	allRe  *regexp.Regexp
	client *irc.Client
}

func (l *listener) listen() {
	for {
		ev := <-l.client.In
		switch ev.Cmd {
		case "CONNECTED":
			log.Printf("Connected.")
		case "PRIVMSG":
			go l.onPrivmsg(ev)
		case "DISCONNECTED":
			log.Fatalf("Disconnected!")
		}
	}
}

func (l *listener) onPrivmsg(ev irc.Event) {
	channel := ev.Args[0]
	line := ev.Args[1]
	for _, m := range l.allRe.FindAllString(line, -1) {
		for _, r := range l.repos {
			if s := r.IssuesRe.FindStringSubmatch(m); s != nil && s[0] == m {
				n := s[2]
				message := fmt.Sprintf("[%s] %s", r.Name, r.IssueURL(n))
				go func() {
					l.client.Out <- irc.Notice{
						Channel: channel,
						Message: message,
					}
				}()
			}
			if s := r.PullsRe.FindStringSubmatch(m); s != nil && s[0] == m {
				n := s[2]
				message := fmt.Sprintf("[%s] %s", r.Name, r.PullURL(n))
				go func() {
					l.client.Out <- irc.Notice{
						Channel: channel,
						Message: message,
					}
				}()
			}
		}
	}
}
<commit_msg>Add more config checks, move them to loadConfig()<commit_after>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"

	"github.com/mvdan/gibot/irc"
	"github.com/mvdan/gibot/site/gitlab"
)

var (
	configPath = flag.String("c", "gibot.json", "path to json config file")
)

var config struct {
	Nick  string   `json:"nick"`
	Chans []string `json:"chans"`
	Repos []struct {
		Name    string   `json:"name"`
		Url     string   `json:"url"`
		Aliases []string `json:"aliases"`
	} `json:"repos"`
}

func main() {
	flag.Parse()
	if err := loadConfig(*configPath); err != nil {
		log.Fatalf("Could not load config: %v", err)
	}
	log.Printf("nick  = %s", config.Nick)
	log.Printf("chans = %s", strings.Join(config.Chans, ", "))

	var repos []gitlab.Repo
	for _, r := range config.Repos {
		aliases := append(r.Aliases, r.Name)
		repos = append(repos, gitlab.NewRepo(r.Name, r.Url, aliases...))
	}
	l := &listener{
		repos: repos,
		allRe: joinRegexes(repos),
	}

	log.Printf("Connecting...")
	c, err := irc.Connect(config.Nick, config.Chans)
	if err != nil {
		log.Fatalf("Could not connect to IRC: %v", err)
	}

	l.client = c
	go l.listen()

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	<-sigc
	log.Printf("Quitting.")
	c.Quit()
}

func loadConfig(p string) error {
	configFile, err := os.Open(p)
	if err != nil {
		return err
	}
	if err := json.NewDecoder(configFile).Decode(&config); err != nil {
		return err
	}
	knownAliases := make(map[string]struct{})
	if config.Nick == "" {
		return fmt.Errorf("no nick specified")
	}
	if len(config.Chans) < 1 {
		return fmt.Errorf("no channels specified")
	}
	for _, r := range config.Repos {
		if r.Name == "" {
			return fmt.Errorf("repo without name")
		}
		if r.Url == "" {
			return fmt.Errorf("repo without url")
		}
		aliases := append(r.Aliases, r.Name)
		for _, a := range aliases {
			if _, e := knownAliases[a]; e {
				return fmt.Errorf("alias '%s' is not unique", a)
			}
			knownAliases[a] = struct{}{}
		}
	}
	return nil
}

func joinRegexes(repos []gitlab.Repo) *regexp.Regexp {
	var all []string
	for _, r := range repos {
		all = append(all, r.IssuesRe.String())
		all = append(all, r.PullsRe.String())
	}
	allRe := regexp.MustCompile(`(` + strings.Join(all, `|`) + `)`)
	allRe.Longest()
	return allRe
}

type listener struct {
	repos  []gitlab.Repo
	allRe  *regexp.Regexp
	client *irc.Client
}

func (l *listener) listen() {
	for {
		ev := <-l.client.In
		switch ev.Cmd {
		case "CONNECTED":
			log.Printf("Connected.")
		case "PRIVMSG":
			go l.onPrivmsg(ev)
		case "DISCONNECTED":
			log.Fatalf("Disconnected!")
		}
	}
}

func (l *listener) onPrivmsg(ev irc.Event) {
	channel := ev.Args[0]
	line := ev.Args[1]
	for _, m := range l.allRe.FindAllString(line, -1) {
		for _, r := range l.repos {
			if s := r.IssuesRe.FindStringSubmatch(m); s != nil && s[0] == m {
				n := s[2]
				message := fmt.Sprintf("[%s] %s", r.Name, r.IssueURL(n))
				go func() {
					l.client.Out <- irc.Notice{
						Channel: channel,
						Message: message,
					}
				}()
			}
			if s := r.PullsRe.FindStringSubmatch(m); s != nil && s[0] == m {
				n := s[2]
				message := fmt.Sprintf("[%s] %s", r.Name, r.PullURL(n))
				go func() {
					l.client.Out <- irc.Notice{
						Channel: channel,
						Message: message,
					}
				}()
			}
		}
	}
}
<|endoftext|>


<commit_before>// elvish is an experimental Unix shell. It tries to incorporate a powerful
// programming language with an extensible, friendly user interface.
package main

import (
	"fmt"
	"os"
	"os/signal"
	"os/user"

	"github.com/elves/elvish/edit"
	"github.com/elves/elvish/errutil"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/parse"
	"github.com/elves/elvish/store"
	"github.com/elves/elvish/util"
)

const (
	sigchSize     = 32
	outChanSize   = 32
	outChanLeader = "▶ "
)

func newEvaluator() *eval.Evaluator {
	ch := make(chan eval.Value, outChanSize)
	go func() {
		for v := range ch {
			fmt.Printf("%s%s\n", outChanLeader, v.Repr())
		}
	}()

	st, err := store.NewStore()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Warning: cannot connect to store:", err)
	}

	ev := eval.NewEvaluator(st)
	ev.SetChanOut(ch)
	return ev
}

func printError(err error) {
	if err != nil {
		if ce, ok := err.(*errutil.ContextualError); ok {
			fmt.Fprint(os.Stderr, ce.Pprint())
		} else {
			fmt.Fprintln(os.Stderr, err.Error())
		}
	}
}

// TODO(xiaq): Currently only the editor deals with signals.
func interact() {
	ev := newEvaluator()

	cmdNum := 0

	username := "???"
	user, err := user.Current()
	if err == nil {
		username = user.Username
	}
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "???"
	}
	rpromptStr := username + "@" + hostname

	sigch := make(chan os.Signal, sigchSize)
	signal.Notify(sigch)

	ed := edit.NewEditor(os.Stdin, ev, sigch)

	for {
		cmdNum++
		name := fmt.Sprintf("<tty %d>", cmdNum)

		prompt := func() string {
			return util.Getwd() + "> "
		}
		rprompt := func() string {
			return rpromptStr
		}

		lr := ed.ReadLine(prompt, rprompt)

		if lr.EOF {
			break
		} else if lr.Err != nil {
			fmt.Println("Editor error:", lr.Err)
			fmt.Println("My pid is", os.Getpid())
		}

		n, pe := parse.Parse(name, lr.Line)
		printError(pe)

		ee := ev.Eval(name, lr.Line, n)
		printError(ee)
	}
}

func script(fname string) {
	ev := newEvaluator()
	err := ev.Source(fname)
	printError(err)
	if err != nil {
		os.Exit(1)
	}
}

var usage = `Usage:
    elvish
    elvish <script>
`

func main() {
	switch len(os.Args) {
	case 1:
		interact()
	case 2:
		script(os.Args[1])
	default:
		fmt.Fprintf(os.Stderr, usage)
		os.Exit(1)
	}
}
<commit_msg>Source ~/.elvish/rc.elv on interact startup.<commit_after>// elvish is an experimental Unix shell. It tries to incorporate a powerful
// programming language with an extensible, friendly user interface.
package main

import (
	"fmt"
	"os"
	"os/signal"
	"os/user"

	"github.com/elves/elvish/edit"
	"github.com/elves/elvish/errutil"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/parse"
	"github.com/elves/elvish/store"
	"github.com/elves/elvish/util"
)

const (
	sigchSize     = 32
	outChanSize   = 32
	outChanLeader = "▶ "
)

func newEvaluator() *eval.Evaluator {
	ch := make(chan eval.Value, outChanSize)
	go func() {
		for v := range ch {
			fmt.Printf("%s%s\n", outChanLeader, v.Repr())
		}
	}()

	st, err := store.NewStore()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Warning: cannot connect to store:", err)
	}

	ev := eval.NewEvaluator(st)
	ev.SetChanOut(ch)
	return ev
}

func printError(err error) {
	if err != nil {
		if ce, ok := err.(*errutil.ContextualError); ok {
			fmt.Fprint(os.Stderr, ce.Pprint())
		} else {
			fmt.Fprintln(os.Stderr, err.Error())
		}
	}
}

// TODO(xiaq): Currently only the editor deals with signals.
func interact() {
	ev := newEvaluator()
	datadir, err := store.EnsureDataDir()
	printError(err)
	if err == nil {
		err := ev.Source(datadir + "/rc.elv")
		if err != nil && !os.IsNotExist(err) {
			printError(err)
		}
	}

	cmdNum := 0

	username := "???"
	user, err := user.Current()
	if err == nil {
		username = user.Username
	}
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "???"
	}
	rpromptStr := username + "@" + hostname

	sigch := make(chan os.Signal, sigchSize)
	signal.Notify(sigch)

	ed := edit.NewEditor(os.Stdin, ev, sigch)

	for {
		cmdNum++
		name := fmt.Sprintf("<tty %d>", cmdNum)

		prompt := func() string {
			return util.Getwd() + "> "
		}
		rprompt := func() string {
			return rpromptStr
		}

		lr := ed.ReadLine(prompt, rprompt)

		if lr.EOF {
			break
		} else if lr.Err != nil {
			fmt.Println("Editor error:", lr.Err)
			fmt.Println("My pid is", os.Getpid())
		}

		n, pe := parse.Parse(name, lr.Line)
		printError(pe)

		ee := ev.Eval(name, lr.Line, n)
		printError(ee)
	}
}

func script(fname string) {
	ev := newEvaluator()
	err := ev.Source(fname)
	printError(err)
	if err != nil {
		os.Exit(1)
	}
}

var usage = `Usage:
    elvish
    elvish <script>
`

func main() {
	switch len(os.Args) {
	case 1:
		interact()
	case 2:
		script(os.Args[1])
	default:
		fmt.Fprintf(os.Stderr, usage)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>/*
Quick and simple monitoring system

Usage:
 jsonmon path.to/config.yml
 jsonmon -v # Prints version to stdout and exits

Environment:
 HOST
  - defaults to localhost
  - the JSON API network interface
 PORT
  - defaults to 3000
  - the JSON API port
 GOMAXPROCS
  - defaults to CPU number + 1
  - number of threads to start
*/
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
	"runtime"
	"strconv"
	"sync"
	"time"
)

// Application version.
const Version = "1.0"

// This one is for internal use.
type ver struct {
	App  string `json:"jsonmon"`
	Go   string `json:"runtime"`
	Os   string `json:"os"`
	Arch string `json:"arch"`
}

var version ver

// Check details.
type Check struct {
	Name   string      `json:"name,omitempty" yaml:"name"`
	Web    string      `json:"web,omitempty" yaml:"web"`
	Shell  string      `json:"shell,omitempty" yaml:"shell"`
	Match  string      `json:"-" yaml:"match"`
	Return int         `json:"-" yaml:"return"`
	Notify interface{} `json:"-" yaml:"notify"`
	Tries  int         `json:"-" yaml:"tries"`
	Repeat int         `json:"-" yaml:"repeat"`
	Failed bool        `json:"failed" yaml:"-"`
	Since  string      `json:"since,omitempty" yaml:"-"`
}

// Global checks list. Need to share it with workers and Web UI.
var checks []Check

// The main loop.
func main() {
	// Parse CLI args.
	usage := "Usage: " + path.Base(os.Args[0]) + " path.to/config.yml\nDocs:  https://github.com/chillum/jsonmon/wiki"
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, usage)
		os.Exit(1)
	}
	// -v for version.
	version.App = Version
	version.Go = runtime.Version()
	version.Os = runtime.GOOS
	version.Arch = runtime.GOARCH
	switch os.Args[1] {
	case "-h":
		fallthrough
	case "-help":
		fallthrough
	case "--help":
		fmt.Fprintln(os.Stderr, usage)
		os.Exit(0)
	case "-v":
		fallthrough
	case "-version":
		fallthrough
	case "--version":
		json, _ := json.MarshalIndent(&version, "", "  ")
		fmt.Println(string(json))
		os.Exit(0)
	}
	// Tune concurrency.
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU() + 1)
	}
	// Read config file or exit with error.
	config, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(3)
	}
	err = yaml.Unmarshal(config, &checks)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", config[0], err)
		os.Exit(3)
	}
	// Run checks.
	var wg sync.WaitGroup
	wg.Add(1)
	for i := range checks {
		go worker(&checks[i])
	}
	// Launch the JSON API.
	host := os.Getenv("HOST")
	port := os.Getenv("PORT")
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "3000"
	}
	http.HandleFunc("/version", getVersion)
	http.HandleFunc("/", getChecks)
	err = http.ListenAndServe(host+":"+port, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
	}
	// Wait forever.
	wg.Wait()
}

// Background worker.
func worker(check *Check) {
	for {
		if check.Repeat == 0 { // Set default timeout.
			check.Repeat = 60
		}
		if check.Tries == 0 { // Default to 1 attempt.
			check.Tries = 1
		}
		if check.Web != "" {
			web(check)
		}
		if check.Shell != "" {
			shell(check)
		}
		time.Sleep(time.Second * time.Duration(check.Repeat))
	}
}

// Shell worker.
func shell(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Shell
	}
	// Execute with shell in N attemps.
	var out []byte
	var err error
	for i := 0; i < check.Tries; i++ {
		out, err = exec.Command("/bin/sh", "-c", check.Shell).CombinedOutput()
		if err == nil {
			if check.Match != "" { // Match regexp.
				var regex *regexp.Regexp
				regex, err = regexp.Compile(check.Match)
				if err == nil && !regex.Match(out) {
					err = errors.New("ERROR: output did not match " + check.Match)
				}
			}
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Fixed: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Failed: "+name, string(out)+err.Error())
		}
	}
}

// Web worker.
func web(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Web
	}
	if check.Return == 0 { // Successful HTTP return code is 200.
		check.Return = 200
	}
	// Get the URL in N attempts.
	var err error
	for i := 0; i < check.Tries; i++ {
		err = fetch(check.Web, check.Match, check.Return)
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FIXED: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FAILED: "+name, err.Error())
		}
	}
}

// The actual HTTP GET.
func fetch(url string, match string, code int) error {
	var err error
	var resp *http.Response
	resp, err = http.Get(url)
	if err == nil && resp.StatusCode != code {
		err = errors.New(url + " returned " + strconv.Itoa(resp.StatusCode))
	} else { // Match regexp.
		if resp != nil && match != "" {
			var regex *regexp.Regexp
			regex, err = regexp.Compile(match)
			if err == nil {
				var body []byte
				body, _ = ioutil.ReadAll(resp.Body)
				if !regex.Match(body) {
					err = errors.New(url + " output did not match " + match)
				}
			}
		}
	}
	if resp != nil {
		resp.Body.Close()
	}
	return err
}

// Logs and mail alerting.
func alert(mail interface{}, subject string, message string) {
	// Log the alerts.
	if message == "" {
		fmt.Println(subject)
	} else {
		fmt.Println(subject + "\n" + message)
	}
	// Mail the alerts.
	if mail != nil {
		// Make the message.
		var rcpt string
		var ok bool
		if rcpt, ok = mail.(string); !ok {
			for i, v := range mail.([]interface{}) {
				if i != 0 {
					rcpt += ", "
				}
				rcpt += v.(string)
			}
		}
		msg := "To: " + rcpt + "\nSubject: " + subject + "\n\n"
		if message != "" {
			msg += message
		}
		msg += "\n.\n"
		// And send it.
		sendmail := exec.Command("/usr/sbin/sendmail", "-t")
		stdin, err := sendmail.StdinPipe()
		err = sendmail.Start()
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR:", err)
		}
		io.WriteString(stdin, msg)
		sendmail.Wait()
	}
}

// Display checks' details.
func getChecks(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" { // Serve for root page only, 404 otherwise.
		http.NotFound(w, r)
		return
	}
	displayJSON(w, &checks)
}

// Display application version.
func getVersion(w http.ResponseWriter, r *http.Request) {
	displayJSON(w, &version)
}

// Output JSON.
func displayJSON(w http.ResponseWriter, data interface{}) {
	json, _ := json.MarshalIndent(&data, "", "  ")
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Write(json)
	fmt.Fprintln(w, "") // Trailing newline.
}
<commit_msg>Server HTTP header<commit_after>/*
Quick and simple monitoring system

Usage:
 jsonmon path.to/config.yml
 jsonmon -v # Prints version to stdout and exits

Environment:
 HOST
  - defaults to localhost
  - the JSON API network interface
 PORT
  - defaults to 3000
  - the JSON API port
 GOMAXPROCS
  - defaults to CPU number + 1
  - number of threads to start
*/
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
	"runtime"
	"strconv"
	"sync"
	"time"
)

// Application version.
const Version = "1.0"

// This one is for internal use.
type ver struct {
	App  string `json:"jsonmon"`
	Go   string `json:"runtime"`
	Os   string `json:"os"`
	Arch string `json:"arch"`
}

var version ver

// Check details.
type Check struct {
	Name   string      `json:"name,omitempty" yaml:"name"`
	Web    string      `json:"web,omitempty" yaml:"web"`
	Shell  string      `json:"shell,omitempty" yaml:"shell"`
	Match  string      `json:"-" yaml:"match"`
	Return int         `json:"-" yaml:"return"`
	Notify interface{} `json:"-" yaml:"notify"`
	Tries  int         `json:"-" yaml:"tries"`
	Repeat int         `json:"-" yaml:"repeat"`
	Failed bool        `json:"failed" yaml:"-"`
	Since  string      `json:"since,omitempty" yaml:"-"`
}

// Global checks list. Need to share it with workers and Web UI.
var checks []Check

// The main loop.
func main() {
	// Parse CLI args.
	usage := "Usage: " + path.Base(os.Args[0]) + " path.to/config.yml\nDocs:  https://github.com/chillum/jsonmon/wiki"
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, usage)
		os.Exit(1)
	}
	// -v for version.
	version.App = Version
	version.Go = runtime.Version()
	version.Os = runtime.GOOS
	version.Arch = runtime.GOARCH
	switch os.Args[1] {
	case "-h":
		fallthrough
	case "-help":
		fallthrough
	case "--help":
		fmt.Fprintln(os.Stderr, usage)
		os.Exit(0)
	case "-v":
		fallthrough
	case "-version":
		fallthrough
	case "--version":
		json, _ := json.MarshalIndent(&version, "", "  ")
		fmt.Println(string(json))
		os.Exit(0)
	}
	// Tune concurrency.
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU() + 1)
	}
	// Read config file or exit with error.
	config, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(3)
	}
	err = yaml.Unmarshal(config, &checks)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", config[0], err)
		os.Exit(3)
	}
	// Run checks.
	var wg sync.WaitGroup
	wg.Add(1)
	for i := range checks {
		go worker(&checks[i])
	}
	// Launch the JSON API.
	host := os.Getenv("HOST")
	port := os.Getenv("PORT")
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "3000"
	}
	http.HandleFunc("/version", getVersion)
	http.HandleFunc("/", getChecks)
	err = http.ListenAndServe(host+":"+port, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
	}
	// Wait forever.
	wg.Wait()
}

// Background worker.
func worker(check *Check) {
	for {
		if check.Repeat == 0 { // Set default timeout.
			check.Repeat = 60
		}
		if check.Tries == 0 { // Default to 1 attempt.
			check.Tries = 1
		}
		if check.Web != "" {
			web(check)
		}
		if check.Shell != "" {
			shell(check)
		}
		time.Sleep(time.Second * time.Duration(check.Repeat))
	}
}

// Shell worker.
func shell(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Shell
	}
	// Execute with shell in N attemps.
	var out []byte
	var err error
	for i := 0; i < check.Tries; i++ {
		out, err = exec.Command("/bin/sh", "-c", check.Shell).CombinedOutput()
		if err == nil {
			if check.Match != "" { // Match regexp.
				var regex *regexp.Regexp
				regex, err = regexp.Compile(check.Match)
				if err == nil && !regex.Match(out) {
					err = errors.New("ERROR: output did not match " + check.Match)
				}
			}
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Fixed: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "Failed: "+name, string(out)+err.Error())
		}
	}
}

// Web worker.
func web(check *Check) {
	// Set check's display name.
	var name string
	if check.Name != "" {
		name = check.Name
	} else {
		name = check.Web
	}
	if check.Return == 0 { // Successful HTTP return code is 200.
		check.Return = 200
	}
	// Get the URL in N attempts.
	var err error
	for i := 0; i < check.Tries; i++ {
		err = fetch(check.Web, check.Match, check.Return)
		if err == nil {
			break
		}
	}
	// Process results.
	if err == nil {
		if check.Failed {
			check.Failed = false
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FIXED: "+name, "")
		}
	} else {
		if !check.Failed {
			check.Failed = true
			check.Since = time.Now().Format(time.RFC3339)
			alert(check.Notify, "FAILED: "+name, err.Error())
		}
	}
}

// The actual HTTP GET.
func fetch(url string, match string, code int) error {
	var err error
	var resp *http.Response
	resp, err = http.Get(url)
	if err == nil && resp.StatusCode != code {
		err = errors.New(url + " returned " + strconv.Itoa(resp.StatusCode))
	} else { // Match regexp.
		if resp != nil && match != "" {
			var regex *regexp.Regexp
			regex, err = regexp.Compile(match)
			if err == nil {
				var body []byte
				body, _ = ioutil.ReadAll(resp.Body)
				if !regex.Match(body) {
					err = errors.New(url + " output did not match " + match)
				}
			}
		}
	}
	if resp != nil {
		resp.Body.Close()
	}
	return err
}

// Logs and mail alerting.
func alert(mail interface{}, subject string, message string) {
	// Log the alerts.
	if message == "" {
		fmt.Println(subject)
	} else {
		fmt.Println(subject + "\n" + message)
	}
	// Mail the alerts.
	if mail != nil {
		// Make the message.
		var rcpt string
		var ok bool
		if rcpt, ok = mail.(string); !ok {
			for i, v := range mail.([]interface{}) {
				if i != 0 {
					rcpt += ", "
				}
				rcpt += v.(string)
			}
		}
		msg := "To: " + rcpt + "\nSubject: " + subject + "\n\n"
		if message != "" {
			msg += message
		}
		msg += "\n.\n"
		// And send it.
		sendmail := exec.Command("/usr/sbin/sendmail", "-t")
		stdin, err := sendmail.StdinPipe()
		err = sendmail.Start()
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR:", err)
		}
		io.WriteString(stdin, msg)
		sendmail.Wait()
	}
}

// Display checks' details.
func getChecks(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" { // Serve for root page only, 404 otherwise.
		w.Header().Set("Server", "jsonmon")
		http.NotFound(w, r)
		return
	}
	displayJSON(w, &checks)
}

// Display application version.
func getVersion(w http.ResponseWriter, r *http.Request) {
	displayJSON(w, &version)
}

// Output JSON.
func displayJSON(w http.ResponseWriter, data interface{}) {
	json, _ := json.MarshalIndent(&data, "", "  ")
	w.Header().Set("Server", "jsonmon")
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Write(json)
	fmt.Fprintln(w, "") // Trailing newline.
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/nacl/box"
)

// Size (in bytes) of the nonce for NaCl box seal/open
const NonceSize = 24

// Size (in bytes) of the key for NaCl box seal/open
const KeySize = 32

func newNonce() ([NonceSize]byte, error) {
	var nonce [NonceSize]byte
	n, err := rand.Read(nonce[:])

	if err != nil {
		return nonce, err
	}

	if n != NonceSize {
		return nonce, errors.New("Not enough bytes read for nonce")
	}

	return nonce, nil
}

type secureReader struct {
	wrappedReader io.Reader
	priv, pub     *[KeySize]byte
}

// decrypt
func (r secureReader) Read(p []byte) (int, error) {
	// make a buffer large enough to handle
	// the overhead associated with an encrypted message
	// (the tag and nonce)
	encryptedMessage := make([]byte, (len(p) + NonceSize + box.Overhead))
	n, err := r.wrappedReader.Read(encryptedMessage)

	if err != nil {
		return n, err
	}

	// strip off the nonce
	// and any extra space at the end of the buffer
	nonceSlice := encryptedMessage[:NonceSize]
	encryptedMessage = encryptedMessage[NonceSize:n]

	// convert the slice to an array
	var nonce [NonceSize]byte
	copy(nonce[:], nonceSlice)

	decryptedMessage, auth := box.Open(nil, encryptedMessage, &nonce, r.pub, r.priv)
	// if authentication failed, output bottom
	if !auth {
		return 0, errors.New("Decrypt error")
	}

	if len(decryptedMessage) > len(p) {
		return 0, errors.New("Decrypt error")
	}

	n = copy(p, decryptedMessage)

	return n, nil
}

type secureWriter struct {
	wrappedWriter io.Writer
	priv, pub     *[KeySize]byte
}

// encrypt
func (w secureWriter) Write(p []byte) (int, error) {
	nonce, err := newNonce()

	if err != nil {
		return 0, err
	}

	encryptedMessage := box.Seal(nil, p, &nonce, w.pub, w.priv)
	// tack the nonce onto the encrypted message
	encWithNonce := append(nonce[:], encryptedMessage...)

	n, err := w.wrappedWriter.Write(encWithNonce)

	if n < len(encWithNonce) {
		return n, errors.New("failed to write complete encrypted message")
	}

	return n, err
}

// NewSecureReader instantiates a new SecureReader
func NewSecureReader(r io.Reader, priv, pub *[KeySize]byte) io.Reader {
	return secureReader{r, priv, pub}
}

// NewSecureWriter instantiates a new SecureWriter
func NewSecureWriter(w io.Writer, priv, pub *[KeySize]byte) io.Writer {
	return secureWriter{w, priv, pub}
}

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func Dial(addr string) (io.ReadWriteCloser, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	// first thing we do is send our public key
	_, err = conn.Write(pub[:])

	if err != nil {
		return nil, err
	}

	// wait for the server's public key
	peerPubSlice := make([]byte, KeySize)
	n, err := conn.Read(peerPubSlice)

	if err != nil {
		return nil, err
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [KeySize]byte
	copy(peerPub[:], peerPubSlice)

	secCon := secureConn{
		NewSecureReader(conn, priv, &peerPub),
		NewSecureWriter(conn, priv, &peerPub),
		conn,
	}

	return secCon, nil
}

// Serve starts a secure echo server on the given listener.
func Serve(l net.Listener) error {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return err
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}

		go handleConnection(conn, pub, priv)
	}
}

func handleConnection(c net.Conn, pub, priv *[32]byte) {
	defer c.Close()
	// first wait for the client's public key
	peerPubSlice := make([]byte, KeySize)
	n, err := c.Read(peerPubSlice)

	if err != nil {
		log.Println(err)
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [KeySize]byte
	copy(peerPub[:], peerPubSlice)

	// then, send our public key
	_, err = c.Write(pub[:])

	if err != nil {
		log.Println(err)
	}

	// now session is "secure"
	secureReader := NewSecureReader(c, priv, &peerPub)
	secureWriter := NewSecureWriter(c, priv, &peerPub)

	// echo
	_, err = io.Copy(secureWriter, secureReader)

	if err != nil {
		log.Println(err)
	}
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(Serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := Dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<commit_msg>add ReadFrom method to SecureWriter in order to prevent the short write error when using io.Copy<commit_after>package main

import (
	"crypto/rand"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/nacl/box"
)

// Size (in bytes) of the nonce for NaCl box seal/open
const NonceSize = 24

// Size (in bytes) of the key for NaCl box seal/open
const KeySize = 32

// Size (in bytes) of total additional overhead added by encryption
const TotalOverhead = NonceSize + box.Overhead

// Size (in bytes) of the max message size supported by this package
const MaxMessageSize = 32 * 1024

func newNonce() ([NonceSize]byte, error) {
	var nonce [NonceSize]byte
	n, err := rand.Read(nonce[:])

	if err != nil {
		return nonce, err
	}

	if n != NonceSize {
		return nonce, errors.New("not enough bytes read for nonce")
	}

	return nonce, nil
}

type secureReader struct {
	r         io.Reader
	priv, pub *[KeySize]byte
}

// Read decrypts a stream encrypted with box.Seal.
// It expects the nonce used to be prepended
// to the ciphertext
func (s secureReader) Read(p []byte) (int, error) {
	// make a buffer large enough to handle
	// the overhead associated with an encrypted message
	// (the tag and nonce)
	enc := make([]byte, (len(p) + TotalOverhead))
	n, err := s.r.Read(enc)

	if err != nil {
		return n, err
	}

	// strip off the nonce
	// and any extra space at the end of the buffer
	nonceSlice := enc[:NonceSize]
	enc = enc[NonceSize:n]

	// convert the slice to an array
	// for use in box.Open
	var nonce [NonceSize]byte
	copy(nonce[:], nonceSlice)

	decrypt, auth := box.Open(nil, enc, &nonce, s.pub, s.priv)
	// if authentication failed, output bottom
	if !auth {
		return 0, errors.New("decrypt error")
	}

	if len(decrypt) > len(p) {
		return 0, errors.New("decrypt error")
	}

	n = copy(p, decrypt)

	return n, nil
}

type secureWriter struct {
	w         io.Writer
	priv, pub *[KeySize]byte
}

// ReadFrom reads from a reader (assumed to be a SecureReader)
// and writes to the SecureWriter.
// This implementation is almost identical to the default
// implementation in io.Copy, but it takes into account
// that the written (encrypted) message is expected to be larger than
// the read (plaintext) message.
func (s secureWriter) ReadFrom(r io.Reader) (int64, error) {
	buf := make([]byte, (MaxMessageSize + TotalOverhead))

	var n int64
	var err error

	for {
		nr, er := r.Read(buf)
		if nr > 0 {
			nw, ew := s.Write(buf[:nr])
			if nw > 0 {
				n += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if (nr + TotalOverhead) != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return n, err
}

// Write encrypts a plaintext stream using box.Seal
func (s secureWriter) Write(p []byte) (int, error) {
	nonce, err := newNonce()

	if err != nil {
		return 0, err
	}

	enc := box.Seal(nil, p, &nonce, s.pub, s.priv)
	// tack the nonce onto the encrypted message
	encWithNonce := append(nonce[:], enc...)

	n, err := s.w.Write(encWithNonce)

	// return an error if the complete message wasn't written
	// this case also fulfills the contract that Write must return
	// an error if n < len(p) since len(encWithNonce) is guaranteed
	// to be greater than len(p)
	if n < len(encWithNonce) {
		return n, errors.New("failed to write complete encrypted message")
	}

	return n, err
}

// NewSecureReader instantiates a new SecureReader
func NewSecureReader(r io.Reader, priv, pub *[KeySize]byte) io.Reader {
	return secureReader{r, priv, pub}
}

// NewSecureWriter instantiates a new SecureWriter
func NewSecureWriter(w io.Writer, priv, pub *[KeySize]byte) io.Writer {
	return secureWriter{w, priv, pub}
}

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func Dial(addr string) (io.ReadWriteCloser, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	// first thing we do is send our public key
	_, err = conn.Write(pub[:])

	if err != nil {
		return nil, err
	}

	// wait for the server's public key
	peerPubSlice := make([]byte, KeySize)
	n, err := conn.Read(peerPubSlice)

	if err != nil {
		return nil, err
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [KeySize]byte
	copy(peerPub[:], peerPubSlice)

	secCon := secureConn{
		NewSecureReader(conn, priv, &peerPub),
		NewSecureWriter(conn, priv, &peerPub),
		conn,
	}

	return secCon, nil
}

// Serve starts a secure echo server on the given listener.
func Serve(l net.Listener) error {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return err
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}

		go handleConnection(conn, pub, priv)
	}
}

func handleConnection(c net.Conn, pub, priv *[32]byte) {
	defer c.Close()
	// first wait for the client's public key
	peerPubSlice := make([]byte, KeySize)
	n, err := c.Read(peerPubSlice)

	if err != nil {
		log.Println(err)
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [KeySize]byte
	copy(peerPub[:], peerPubSlice)

	// then, send our public key
	_, err = c.Write(pub[:])

	if err != nil {
		log.Println(err)
	}

	// now session is "secure"
	sr := NewSecureReader(c, priv, &peerPub)
	sw := NewSecureWriter(c, priv, &peerPub)

	// echo
	_, err = io.Copy(sw, sr)

	if err != nil {
		log.Println(err)
	}
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(Serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := Dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<|endoftext|>