<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
)

type protoImplT func(bool, bool, string, string, []string) error

const logsEnabled = true
const usage = "Usage: goon -proto <version> [options] -- <program> [<arg>...]"

var protoFlag = flag.String("proto", "", "protocol version (one of: 0.0)")
var inFlag  = flag.Bool("in", false, "whether stdin is used")
var outFlag = flag.Bool("out", false, "whether stdout is preserved or discarded")
var errFlag = flag.String("err", "nil", "redirection or supression of stderr")
var dirFlag = flag.String("dir", ".", "working directory for the spawned process")

func main() {
	flag.Parse()
	args := flag.Args()

	validateOptsAndArgs(*protoFlag, args)

	/* Run external program and block until it terminates */
	err := findProtocolImpl(*protoFlag)(*inFlag, *outFlag, *errFlag, *dirFlag, args)

	/* Determine the exit status */
	if err != nil {
		//fmt.Printf("%#v\n", err)
		os.Exit(getExitStatus(err))
	}
}

func validateOptsAndArgs(protoFlag string, args []string) {
	if protoFlag == "" {
		die_usage("Please specify the protocol version.")
	}

	if len(args) < 1 {
		die_usage("Not enough arguments.")
	}

	logger.Printf("Flag values:\n  proto: %v\n  in: %v\n  out: %v\n  err: %v\n  dir: %v\nArgs: %v\n",
				  protoFlag, *inFlag, *outFlag, *errFlag, *dirFlag, args)
}

func findProtocolImpl(flag string) (impl protoImplT) {
	switch flag {
	case "0.0":
		impl = proto_0_0
	default:
		reason := fmt.Sprintf("Unsupported protocol version: %v", flag)
		die(reason)
	}
	return
}
<commit_msg>Disable logging<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
)

type protoImplT func(bool, bool, string, string, []string) error

const logsEnabled = false
const usage = "Usage: goon -proto <version> [options] -- <program> [<arg>...]"

var protoFlag = flag.String("proto", "", "protocol version (one of: 0.0)")
var inFlag  = flag.Bool("in", false, "whether stdin is used")
var outFlag = flag.Bool("out", false, "whether stdout is preserved or discarded")
var errFlag = flag.String("err", "nil", "redirection or supression of stderr")
var dirFlag = flag.String("dir", ".", "working directory for the spawned process")

func main() {
	flag.Parse()
	args := flag.Args()

	validateOptsAndArgs(*protoFlag, args)

	/* Run external program and block until it terminates */
	err := findProtocolImpl(*protoFlag)(*inFlag, *outFlag, *errFlag, *dirFlag, args)

	/* Determine the exit status */
	if err != nil {
		//fmt.Printf("%#v\n", err)
		os.Exit(getExitStatus(err))
	}
}

func validateOptsAndArgs(protoFlag string, args []string) {
	if protoFlag == "" {
		die_usage("Please specify the protocol version.")
	}

	if len(args) < 1 {
		die_usage("Not enough arguments.")
	}

	logger.Printf("Flag values:\n  proto: %v\n  in: %v\n  out: %v\n  err: %v\n  dir: %v\nArgs: %v\n",
				  protoFlag, *inFlag, *outFlag, *errFlag, *dirFlag, args)
}

func findProtocolImpl(flag string) (impl protoImplT) {
	switch flag {
	case "0.0":
		impl = proto_0_0
	default:
		reason := fmt.Sprintf("Unsupported protocol version: %v", flag)
		die(reason)
	}
	return
}
<|endoftext|>


<commit_before>// get-headers prints the headers from GET requesting a URL
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	netURL "net/url"
	"os"
	"time"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

const usage = `Usage of get-headers:

get-headers [opts] <url>...
       	Gets the URLs and prints their headers alphabetically.
       	Repeated headers are printed with an asterisk.
  -g	Shortcut for -gzip
  -gzip
       	Enable GZIP compression

get-headers (-h|--help)
	Print this help message.
`

var (
	gzip = flag.Bool("gzip", false, "Enable GZIP compression")
)

func init() {
	flag.BoolVar(gzip, "g", false, "Shortcut for -gzip")
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, usage)
	}
	flag.Parse()
	if len(flag.Args()) < 1 {
		flag.Usage()
		os.Exit(2)
	}
}

// Sentinal error to let us know if we're ignoring a redirect
var errRedirect = errors.New("redirected")

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return errRedirect
}

func die(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "get-time: %v\n", err)
		os.Exit(1)
	}
}

func main() {
	transport := &http.Transport{
		DisableCompression: *gzip,
	}
	client := http.Client{
		CheckRedirect: checkRedirect,
		Transport:     transport,
	}
	for _, url := range flag.Args() {
		var n int64
		start := time.Now()

		resp, err := client.Get(url)

		// Ignore the error if it's just our errRedirect
		switch urlErr, ok := err.(*netURL.Error); {
		case err == nil:
			// Copying to /dev/null just to make sure this is real
			n, err = io.Copy(ioutil.Discard, resp.Body)
			die(err)
		case ok && urlErr.Err == errRedirect:
		default:
			die(err)
		}
		duration := time.Since(start)
		die(resp.Body.Close())

		fmt.Println("GET", url)
		fmt.Println(resp.Proto, resp.Status, "\n")
		fmt.Println("Time           ", humanizeDuration(duration))
		if n != 0 {
			fmt.Println("Content length ", humanizeByteSize(n))
			bps := int64(float64(n) / duration.Seconds())
			fmt.Printf("Speed           %s/s\n", humanizeByteSize(bps))
		}
		fmt.Println()
		fmt.Println(prettyprint.ResponseHeader(resp.Header))

	}
}
<commit_msg>Set accept-headers for GZIP<commit_after>// get-headers prints the headers from GET requesting a URL
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	netURL "net/url"
	"os"
	"time"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

const usage = `Usage of get-headers:

get-headers [opts] <url>...
       	Gets the URLs and prints their headers alphabetically.
       	Repeated headers are printed with an asterisk.
  -g	Shortcut for -gzip
  -gzip
       	Enable GZIP compression

get-headers (-h|--help)
	Print this help message.
`

var (
	gzip = flag.Bool("gzip", false, "Enable GZIP compression")
)

func init() {
	flag.BoolVar(gzip, "g", false, "Shortcut for -gzip")
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, usage)
	}
	flag.Parse()
	if len(flag.Args()) < 1 {
		flag.Usage()
		os.Exit(2)
	}
}

// Sentinal error to let us know if we're ignoring a redirect
var errRedirect = errors.New("redirected")

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return errRedirect
}

func die(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "get-time: %v\n", err)
		os.Exit(1)
	}
}

func main() {
	transport := &http.Transport{
		DisableCompression: *gzip,
	}
	client := http.Client{
		CheckRedirect: checkRedirect,
		Transport:     transport,
	}
	for _, url := range flag.Args() {
		var n int64
		start := time.Now()

		req, err := http.NewRequest("GET", url, nil)
		die(err)
		if *gzip {
			req.Header = map[string][]string{
				"Accept-Encoding": {"gzip, deflate"},
			}
		}
		resp, err := client.Do(req)

		// Ignore the error if it's just our errRedirect
		switch urlErr, ok := err.(*netURL.Error); {
		case err == nil:
			// Copying to /dev/null just to make sure this is real
			n, err = io.Copy(ioutil.Discard, resp.Body)
			die(err)
		case ok && urlErr.Err == errRedirect:
		default:
			die(err)
		}
		duration := time.Since(start)
		die(resp.Body.Close())

		fmt.Println("GET", url)
		fmt.Println(resp.Proto, resp.Status, "\n")
		fmt.Println("Time           ", humanizeDuration(duration))
		if n != 0 {
			fmt.Println("Content length ", humanizeByteSize(n))
			bps := int64(float64(n) / duration.Seconds())
			fmt.Printf("Speed           %s/s\n", humanizeByteSize(bps))
		}
		fmt.Println()
		fmt.Println(prettyprint.ResponseHeader(resp.Header))

	}
}
<|endoftext|>


<commit_before>package main // import "github.com/CenturyLinkLabs/watchtower"

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/CenturyLinkLabs/watchtower/actions"
	"github.com/CenturyLinkLabs/watchtower/container"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
)

var (
	wg           sync.WaitGroup
	client       container.Client
	pollInterval time.Duration
	cleanup      bool
)

func init() {
	log.SetLevel(log.InfoLevel)
}

func main() {
	rootCertPath := "/etc/ssl/docker"

	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "watchtower"
	app.Usage = "Automatically update running Docker containers"
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.IntFlag{
			Name:  "interval, i",
			Usage: "poll interval (in seconds)",
			Value: 300,
		},
		cli.BoolFlag{
			Name:  "no-pull",
			Usage: "do not pull new images",
		},
		cli.BoolFlag{
			Name:  "cleanup",
			Usage: "remove old images after updating",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	pollInterval = time.Duration(c.Int("interval")) * time.Second
	cleanup = c.GlobalBool("cleanup")

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}

	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	handleSignals()
	return nil
}

func start(c *cli.Context) {
	names := c.Args()

	if err := actions.CheckPrereqs(client, cleanup); err != nil {
		log.Fatal(err)
	}

	for {
		wg.Add(1)
		if err := actions.Update(client, names, cleanup); err != nil {
			fmt.Println(err)
		}
		wg.Done()

		time.Sleep(pollInterval)
	}
}

func handleSignals() {
	// Graceful shut-down on SIGINT/SIGTERM
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)

	go func() {
		<-c
		wg.Wait()
		os.Exit(1)
	}()
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<commit_msg>Setup using env vars as well. Add no retsart option<commit_after>package main // import "github.com/CenturyLinkLabs/watchtower"

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/CenturyLinkLabs/watchtower/actions"
	"github.com/CenturyLinkLabs/watchtower/container"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
)

var (
	wg           sync.WaitGroup
	client       container.Client
	pollInterval time.Duration
	cleanup      bool
	noRestart    bool
)

func init() {
	log.SetLevel(log.InfoLevel)
}

func main() {
	rootCertPath := "/etc/ssl/docker"

	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "watchtower"
	app.Usage = "Automatically update running Docker containers"
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.IntFlag{
			Name:   "interval, i",
			Usage:  "poll interval (in seconds)",
			Value:  300,
			EnvVar: "WATCHTOWER_POLL_INTERVAL",
		},
		cli.BoolFlag{
			Name:   "no-pull",
			Usage:  "do not pull new images",
			EnvVar: "WATCHTOWER_NO_PULL",
		},
		cli.BoolFlag{
			Name:   "no-restart",
			Usage:  "do not restart containers",
			EnvVar: "WATCHTOWER_NO_PULL",
		},
		cli.BoolFlag{
			Name:   "cleanup",
			Usage:  "remove old images after updating",
			EnvVar: "WATCHTOWER_CLEANUP",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	pollInterval = time.Duration(c.Int("interval")) * time.Second
	cleanup = c.GlobalBool("cleanup")
	noRestart = c.GlobalBool("no-restart")

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}

	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	handleSignals()
	return nil
}

func start(c *cli.Context) {
	names := c.Args()

	if err := actions.CheckPrereqs(client, cleanup); err != nil {
		log.Fatal(err)
	}

	for {
		wg.Add(1)
		if err := actions.Update(client, names, cleanup, noRestart); err != nil {
			fmt.Println(err)
		}
		wg.Done()

		time.Sleep(pollInterval)
	}
}

func handleSignals() {
	// Graceful shut-down on SIGINT/SIGTERM
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)

	go func() {
		<-c
		wg.Wait()
		os.Exit(1)
	}()
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<|endoftext|>


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

package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"

	"github.com/line/line-bot-sdk-go/linebot"
)

var bot *linebot.Client

func main() {
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	events, err := bot.ParseRequest(r)

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {

			switch message := event.Message.(type) {
			case *linebot.TextMessage:

				msg = message.Text
				var keyword [3]string
				keyword["罐罐"] = "肚子裡永遠少一罐"
				keyword["娜娜"] = "叫我嗎？叫一次一百萬"
				keyword["天氣"] = "不管天氣好壞,我都不想離開被窩"

				for index,element := range keyword {
					match, _ := regexp.MatchString(index, msg)
					if( match ){
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(element)).Do(); err != nil {
							log.Print(err)
						}

					} else {
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(msg)).Do(); err != nil {
							log.Print(err)
						}
					}
				}

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

package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"

	"github.com/line/line-bot-sdk-go/linebot"
)

var bot *linebot.Client

func main() {
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	events, err := bot.ParseRequest(r)

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {

			switch message := event.Message.(type) {
			case *linebot.TextMessage:

				msg := message.Text
				var keyword [3]string
				keyword["罐罐"] = "肚子裡永遠少一罐"
				keyword["娜娜"] = "叫我嗎？叫一次一百萬"
				keyword["天氣"] = "不管天氣好壞,我都不想離開被窩"

				for index,element := range keyword {
					match, _ := regexp.MatchString(index, msg)
					if( match ){
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(element)).Do(); err != nil {
							log.Print(err)
						}

					} else {
						if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(msg)).Do(); err != nil {
							log.Print(err)
						}
					}
				}

			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
	"unicode"
)

const (
	CmdDisplayMax = 32
	// max length of /proc/$PID/comm
	CommMax = 16
	// from ps_mem - average error due to truncation in the kernel
	// pss calculations
	PssAdjust = .5
	pageSize  = 4096
	mapDetailLen = len("Size:                  4 kB")

	usage = `Usage: %s [OPTION...]
Simple, accurate RAM and swap reporting.

Options:
`
)

var (
	memProfile string
	cpuProfile string
)

// store info about a command (group of processes), similar to how
// ps_mem works.
type CmdMemInfo struct {
	PIDs    []int
	Name    string
	Pss     float64
	Shared  uint64
	Private uint64
	Swapped uint64
}

type MapInfo struct {
	Inode uint64
	Name  string
}

// mapLine is a line from /proc/$PID/maps, or one of the same header
// lines from smaps.
func NewMapInfo(mapLine []byte) MapInfo {
	var mi MapInfo
	var err error
	pieces := splitSpaces(mapLine)
	if len(pieces) == 6 {
		mi.Name = string(pieces[5])
	}
	mi.Inode, err = ParseUint(pieces[4], 10, 64)
	if err != nil {
		panic(fmt.Sprintf("NewMapInfo: Atoi(%s): %s (%s)",
			string(pieces[4]), err, string(mapLine)))
	}
	return mi
}

func (mi MapInfo) IsAnon() bool {
	return mi.Inode == 0
}

// isDigit returns true if the rune d represents an ascii digit
// between 0 and 9, inclusive.
func isDigit(d uint8) bool {
	return d >= '0' && d <= '9'
}

// pidList returns a list of the process-IDs of every currently
// running process on the local system.
func pidList() ([]int, error) {
	procLs, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, fmt.Errorf("ReadDir(/proc): %s", err)
	}

	pids := make([]int, 0, len(procLs))
	for _, pInfo := range procLs {
		if !isDigit(pInfo.Name()[0]) || !pInfo.IsDir() {
			continue
		}
		pidInt, err := strconv.Atoi(pInfo.Name())
		if err != nil {
			return nil, fmt.Errorf("Atoi(%s): %s", pInfo.Name(), err)
		}
		pids = append(pids, pidInt)
	}
	return pids, nil
}

// procName gets the process name for a worker.  It first checks the
// value of /proc/$PID/cmdline.  If setproctitle(3) has been called,
// it will use this.  Otherwise it uses the value of
// path.Base(/proc/$PID/exe), which has info on whether the executable
// has changed since the process was exec'ed.
func procName(pid int) (string, error) {
	p, err := os.Readlink(fmt.Sprintf("/proc/%d/exe", pid))
	// this would return an error if the PID doesn't
	// exist, or if the PID refers to a kernel thread.
	if err != nil {
		return "", nil
	}
	// cmdline is the null separated list of command line
	// arguments for the process, unless setproctitle(3)
	// has been called, in which case it is the
	argsB, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
	if err != nil {
		return "", fmt.Errorf("ReadFile(%s): %s", fmt.Sprintf("/proc/%d/cmdline", pid), err)
	}
	args := strings.Split(string(argsB), "\000")
	n := args[0]

	nTrunc := n
	if len(n) > CommMax {
		nTrunc = n[:CommMax]
	}
	if strings.HasPrefix(p, nTrunc) {
		n = path.Base(p)
	}
	return n, nil
}

// splitSpaces returns a slice of byte slices which are the space
// delimited words from the original byte slice.  Unlike
// strings.Split($X, " "), runs of multiple spaces in a row are
// discarded.
func splitSpaces(b []byte) [][]byte {
	res := make([][]byte, 0, 6) // 6 is empirically derived
	start, i := 0, 0
	lenB := len(b)
	for i = 0; i < lenB; i++ {
		// fast forward past any spaces
		for i < lenB-1 && unicode.IsSpace(rune(b[i])) {
			i++
			start = i
		}
		for i < lenB-1 && !unicode.IsSpace(rune(b[i])) {
			i++
		}
		if i > start {
			// we sometimes have to rewind
			if i < lenB-1 && unicode.IsSpace(rune(b[i])) {
				i--
			}
			res = append(res, b[start:i+1])
			start = i + 1
		}
	}
	return res
}

// procMem returns the amount of Pss, shared, and swapped out memory
// used.  The swapped out amount refers to anonymous pages only.
func procMem(pid int) (pss float64, shared, priv, swap uint64, err error) {
	fPath := fmt.Sprintf("/proc/%d/smaps", pid)
	f, err := os.Open(fPath)
	if err != nil {
		err = fmt.Errorf("os.Open(%s): %s", fPath, err)
		return
	}
	//var curr MapInfo
	r := bufio.NewReaderSize(f, pageSize)
	for {
		var l []byte
		var isPrefix bool
		l, isPrefix, err = r.ReadLine()
		// this should never happen, so take the easy way out.
		if isPrefix {
			err = fmt.Errorf("ReadLine(%s): isPrefix", fPath)
		}
		if err != nil {
			// if we've got EOF, then we're simply done
			// processing smaps.
			if err == io.EOF {
				err = nil
				break
			}
			// otherwise error out
			err = fmt.Errorf("ReadLine(%s): %s", fPath, err)
			return
		}

		if len(l) != mapDetailLen {
			//curr = NewMapInfo(l)
			continue
		}
		pieces := splitSpaces(l)
		ty := string(pieces[0])
		var v uint64
		switch ty {
		case "Pss:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			pss += float64(v) + PssAdjust
		case "Shared_Clean:", "Shared_Dirty:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			shared += v
		case "Private_Clean:", "Private_Dirty:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			priv += v
		case "Swap:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			swap += v
		}
	}
	return
}

func worker(pidRequest chan int, wg *sync.WaitGroup, result chan *CmdMemInfo) {
	for pid := range pidRequest {
		var err error
		cmi := new(CmdMemInfo)

		cmi.PIDs = []int{pid}
		cmi.Name, err = procName(pid)
		if err != nil {
			log.Printf("procName(%d): %s", pid, err)
			wg.Done()
			continue
		} else if cmi.Name == "" {
			// XXX: This happens with kernel
			// threads. maybe warn? idk.
			wg.Done()
			continue
		}

		cmi.Pss, cmi.Shared, cmi.Private, cmi.Swapped, err = procMem(pid)
		if err != nil {
			log.Printf("procMem(%d): %s", pid, err)
			wg.Done()
			continue
		}

		result <- cmi
		wg.Done()
	}
}

type byTotal []*CmdMemInfo

func (c byTotal) Len() int           { return len(c) }
func (c byTotal) Less(i, j int) bool { return c[i].Pss < c[j].Pss }
func (c byTotal) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")

	flag.Parse()
}

// startProfiling enables memory and/or CPU profiling if the
// appropriate command line flags have been set.
func startProfiling() {
	var err error
	// if we've passed in filenames to dump profiling data too,
	// start collecting profiling data.
	if memProfile != "" {
		runtime.MemProfileRate = 1
	}
	if cpuProfile != "" {
		var f *os.File
		if f, err = os.Create(cpuProfile); err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}
}

func stopProfiling() {
	if memProfile != "" {
		runtime.GC()
		f, err := os.Create(memProfile)
		if err != nil {
			log.Println(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
	if cpuProfile != "" {
		pprof.StopCPUProfile()
		cpuProfile = ""
	}
}

func main() {
	startProfiling()
	defer stopProfiling()

	// need to be root to read map info for other user's
	// processes.
	if os.Geteuid() != 0 {
		fmt.Printf("FATAL: root required.\n")
		return
	}

	pids, err := pidList()
	if err != nil {
		log.Printf("pidList: %s", err)
		return
	}

	var wg sync.WaitGroup
	work := make(chan int, len(pids))
	result := make(chan *CmdMemInfo, len(pids))

	// give us as much parallelism as possible
	nCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(nCPU)
	for i := 0; i < nCPU; i++ {
		go worker(work, &wg, result)
	}

	wg.Add(len(pids))
	for _, pid := range pids {
		work <- pid
	}
	wg.Wait()

	// aggregate similar processes by command name.
	cmdMap := map[string]*CmdMemInfo{}
loop:
	for {
		// this only works correctly because we a channel
		// where the buffer size >= the number of potential
		// results.
		select {
		case c := <-result:
			n := c.Name
			if _, ok := cmdMap[n]; !ok {
				cmdMap[n] = c
				continue
			}
			cmdMap[n].PIDs = append(cmdMap[n].PIDs, c.PIDs...)
			cmdMap[n].Pss += c.Pss
			cmdMap[n].Private += c.Private
			cmdMap[n].Shared += c.Shared
			cmdMap[n].Swapped += c.Swapped
		default:
			break loop
		}
	}

	// extract map values to a slice so we can sort them
	cmds := make([]*CmdMemInfo, 0, len(cmdMap))
	for _, c := range cmdMap {
		cmds = append(cmds, c)
	}
	sort.Sort(byTotal(cmds))

	// keep track of total RAM and swap usage
	var totPss, totSwap float64

	fmt.Printf("%10s%10s%10s\t%s\n", "MB RAM", "PRIVATE", "SWAPPED", "PROCESS (COUNT)")
	for _, c := range cmds {
		n := c.Name
		if len(n) > CmdDisplayMax {
			if n[0] == '[' {
				n = n[:strings.IndexRune(n, ']')+1]
			} else {
				n = n[:CmdDisplayMax]
			}
		}
		s := ""
		if c.Swapped > 0 {
			swap := float64(c.Swapped) / 1024.
			totSwap += swap
			s = fmt.Sprintf("%10.1f", swap)
		}
		pss := float64(c.Pss) / 1024.
		fmt.Printf("%10.1f%10.1f%10s\t%s (%d)\n", pss, float64(c.Private)/1024, s, n, len(c.PIDs))
		totPss += pss
	}
	fmt.Printf("#%9.1f%20.1f\tTOTAL USED BY PROCESSES\n", totPss, totSwap)
}
<commit_msg>splitSpaces: check if byte == ' ' directly<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
)

const (
	CmdDisplayMax = 32
	// max length of /proc/$PID/comm
	CommMax = 16
	// from ps_mem - average error due to truncation in the kernel
	// pss calculations
	PssAdjust = .5
	pageSize  = 4096
	mapDetailLen = len("Size:                  4 kB")

	usage = `Usage: %s [OPTION...]
Simple, accurate RAM and swap reporting.

Options:
`
)

var (
	memProfile string
	cpuProfile string
)

// store info about a command (group of processes), similar to how
// ps_mem works.
type CmdMemInfo struct {
	PIDs    []int
	Name    string
	Pss     float64
	Shared  uint64
	Private uint64
	Swapped uint64
}

type MapInfo struct {
	Inode uint64
	Name  string
}

// mapLine is a line from /proc/$PID/maps, or one of the same header
// lines from smaps.
func NewMapInfo(mapLine []byte) MapInfo {
	var mi MapInfo
	var err error
	pieces := splitSpaces(mapLine)
	if len(pieces) == 6 {
		mi.Name = string(pieces[5])
	}
	mi.Inode, err = ParseUint(pieces[4], 10, 64)
	if err != nil {
		panic(fmt.Sprintf("NewMapInfo: Atoi(%s): %s (%s)",
			string(pieces[4]), err, string(mapLine)))
	}
	return mi
}

func (mi MapInfo) IsAnon() bool {
	return mi.Inode == 0
}

// isDigit returns true if the rune d represents an ascii digit
// between 0 and 9, inclusive.
func isDigit(d uint8) bool {
	return d >= '0' && d <= '9'
}

// pidList returns a list of the process-IDs of every currently
// running process on the local system.
func pidList() ([]int, error) {
	procLs, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, fmt.Errorf("ReadDir(/proc): %s", err)
	}

	pids := make([]int, 0, len(procLs))
	for _, pInfo := range procLs {
		if !isDigit(pInfo.Name()[0]) || !pInfo.IsDir() {
			continue
		}
		pidInt, err := strconv.Atoi(pInfo.Name())
		if err != nil {
			return nil, fmt.Errorf("Atoi(%s): %s", pInfo.Name(), err)
		}
		pids = append(pids, pidInt)
	}
	return pids, nil
}

// procName gets the process name for a worker.  It first checks the
// value of /proc/$PID/cmdline.  If setproctitle(3) has been called,
// it will use this.  Otherwise it uses the value of
// path.Base(/proc/$PID/exe), which has info on whether the executable
// has changed since the process was exec'ed.
func procName(pid int) (string, error) {
	p, err := os.Readlink(fmt.Sprintf("/proc/%d/exe", pid))
	// this would return an error if the PID doesn't
	// exist, or if the PID refers to a kernel thread.
	if err != nil {
		return "", nil
	}
	// cmdline is the null separated list of command line
	// arguments for the process, unless setproctitle(3)
	// has been called, in which case it is the
	argsB, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
	if err != nil {
		return "", fmt.Errorf("ReadFile(%s): %s", fmt.Sprintf("/proc/%d/cmdline", pid), err)
	}
	args := strings.Split(string(argsB), "\000")
	n := args[0]

	nTrunc := n
	if len(n) > CommMax {
		nTrunc = n[:CommMax]
	}
	if strings.HasPrefix(p, nTrunc) {
		n = path.Base(p)
	}
	return n, nil
}

// splitSpaces returns a slice of byte slices which are the space
// delimited words from the original byte slice.  Unlike
// strings.Split($X, " "), runs of multiple spaces in a row are
// discarded.  NOTE WELL: this only checks for spaces (' '), other
// unicode whitespace isn't supported.
func splitSpaces(b []byte) [][]byte {
	res := make([][]byte, 0, 6) // 6 is empirically derived
	start, i := 0, 0
	lenB := len(b)
	for i = 0; i < lenB; i++ {
		// fast forward past any spaces
		for i < lenB-1 && b[i] == ' ' {
			i++
			start = i
		}
		for i < lenB-1 && b[i] != ' ' {
			i++
		}
		if i > start {
			// we sometimes have to rewind
			if i < lenB-1 && b[i] == ' ' {
				i--
			}
			res = append(res, b[start:i+1])
			start = i + 1
		}
	}
	return res
}

// procMem returns the amount of Pss, shared, and swapped out memory
// used.  The swapped out amount refers to anonymous pages only.
func procMem(pid int) (pss float64, shared, priv, swap uint64, err error) {
	fPath := fmt.Sprintf("/proc/%d/smaps", pid)
	f, err := os.Open(fPath)
	if err != nil {
		err = fmt.Errorf("os.Open(%s): %s", fPath, err)
		return
	}
	//var curr MapInfo
	r := bufio.NewReaderSize(f, pageSize)
	for {
		var l []byte
		var isPrefix bool
		l, isPrefix, err = r.ReadLine()
		// this should never happen, so take the easy way out.
		if isPrefix {
			err = fmt.Errorf("ReadLine(%s): isPrefix", fPath)
		}
		if err != nil {
			// if we've got EOF, then we're simply done
			// processing smaps.
			if err == io.EOF {
				err = nil
				break
			}
			// otherwise error out
			err = fmt.Errorf("ReadLine(%s): %s", fPath, err)
			return
		}

		if len(l) != mapDetailLen {
			//curr = NewMapInfo(l)
			continue
		}
		pieces := splitSpaces(l)
		ty := string(pieces[0])
		var v uint64
		switch ty {
		case "Pss:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			pss += float64(v) + PssAdjust
		case "Shared_Clean:", "Shared_Dirty:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			shared += v
		case "Private_Clean:", "Private_Dirty:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			priv += v
		case "Swap:":
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			swap += v
		}
	}
	return
}

func worker(pidRequest chan int, wg *sync.WaitGroup, result chan *CmdMemInfo) {
	for pid := range pidRequest {
		var err error
		cmi := new(CmdMemInfo)

		cmi.PIDs = []int{pid}
		cmi.Name, err = procName(pid)
		if err != nil {
			log.Printf("procName(%d): %s", pid, err)
			wg.Done()
			continue
		} else if cmi.Name == "" {
			// XXX: This happens with kernel
			// threads. maybe warn? idk.
			wg.Done()
			continue
		}

		cmi.Pss, cmi.Shared, cmi.Private, cmi.Swapped, err = procMem(pid)
		if err != nil {
			log.Printf("procMem(%d): %s", pid, err)
			wg.Done()
			continue
		}

		result <- cmi
		wg.Done()
	}
}

type byTotal []*CmdMemInfo

func (c byTotal) Len() int           { return len(c) }
func (c byTotal) Less(i, j int) bool { return c[i].Pss < c[j].Pss }
func (c byTotal) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")

	flag.Parse()
}

// startProfiling enables memory and/or CPU profiling if the
// appropriate command line flags have been set.
func startProfiling() {
	var err error
	// if we've passed in filenames to dump profiling data too,
	// start collecting profiling data.
	if memProfile != "" {
		runtime.MemProfileRate = 1
	}
	if cpuProfile != "" {
		var f *os.File
		if f, err = os.Create(cpuProfile); err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}
}

func stopProfiling() {
	if memProfile != "" {
		runtime.GC()
		f, err := os.Create(memProfile)
		if err != nil {
			log.Println(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
	if cpuProfile != "" {
		pprof.StopCPUProfile()
		cpuProfile = ""
	}
}

func main() {
	startProfiling()
	defer stopProfiling()

	// need to be root to read map info for other user's
	// processes.
	if os.Geteuid() != 0 {
		fmt.Printf("FATAL: root required.\n")
		return
	}

	pids, err := pidList()
	if err != nil {
		log.Printf("pidList: %s", err)
		return
	}

	var wg sync.WaitGroup
	work := make(chan int, len(pids))
	result := make(chan *CmdMemInfo, len(pids))

	// give us as much parallelism as possible
	nCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(nCPU)
	for i := 0; i < nCPU; i++ {
		go worker(work, &wg, result)
	}

	wg.Add(len(pids))
	for _, pid := range pids {
		work <- pid
	}
	wg.Wait()

	// aggregate similar processes by command name.
	cmdMap := map[string]*CmdMemInfo{}
loop:
	for {
		// this only works correctly because we a channel
		// where the buffer size >= the number of potential
		// results.
		select {
		case c := <-result:
			n := c.Name
			if _, ok := cmdMap[n]; !ok {
				cmdMap[n] = c
				continue
			}
			cmdMap[n].PIDs = append(cmdMap[n].PIDs, c.PIDs...)
			cmdMap[n].Pss += c.Pss
			cmdMap[n].Private += c.Private
			cmdMap[n].Shared += c.Shared
			cmdMap[n].Swapped += c.Swapped
		default:
			break loop
		}
	}

	// extract map values to a slice so we can sort them
	cmds := make([]*CmdMemInfo, 0, len(cmdMap))
	for _, c := range cmdMap {
		cmds = append(cmds, c)
	}
	sort.Sort(byTotal(cmds))

	// keep track of total RAM and swap usage
	var totPss, totSwap float64

	fmt.Printf("%10s%10s%10s\t%s\n", "MB RAM", "PRIVATE", "SWAPPED", "PROCESS (COUNT)")
	for _, c := range cmds {
		n := c.Name
		if len(n) > CmdDisplayMax {
			if n[0] == '[' {
				n = n[:strings.IndexRune(n, ']')+1]
			} else {
				n = n[:CmdDisplayMax]
			}
		}
		s := ""
		if c.Swapped > 0 {
			swap := float64(c.Swapped) / 1024.
			totSwap += swap
			s = fmt.Sprintf("%10.1f", swap)
		}
		pss := float64(c.Pss) / 1024.
		fmt.Printf("%10.1f%10.1f%10s\t%s (%d)\n", pss, float64(c.Private)/1024, s, n, len(c.PIDs))
		totPss += pss
	}
	fmt.Printf("#%9.1f%20.1f\tTOTAL USED BY PROCESSES\n", totPss, totSwap)
}
<|endoftext|>


<commit_before>// Package main defines a command line interface for the sqlboiler package
package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/kat-co/vala"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/vattle/sqlboiler/bdb/drivers"
	"github.com/vattle/sqlboiler/boilingcore"
)

const sqlBoilerVersion = "2.1.8"

var (
	cmdState  *boilingcore.State
	cmdConfig *boilingcore.Config
)

func main() {
	var err error

	// Too much happens between here and cobra's argument handling, for
	// something so simple just do it immediately.
	for _, arg := range os.Args {
		if arg == "--version" {
			fmt.Println("SQLBoiler v" + sqlBoilerVersion)
			return
		}
	}

	viper.SetConfigName("sqlboiler")

	configHome := os.Getenv("XDG_CONFIG_HOME")
	homePath := os.Getenv("HOME")
	wd, err := os.Getwd()
	if err != nil {
		wd = "./"
	}

	configPaths := []string{wd}
	if len(configHome) > 0 {
		configPaths = append(configPaths, filepath.Join(configHome, "sqlboiler"))
	} else {
		configPaths = append(configPaths, filepath.Join(homePath, ".config/sqlboiler"))
	}

	for _, p := range configPaths {
		viper.AddConfigPath(p)
	}

	// Ignore errors here, fallback to other validation methods.
	// Users can use environment variables if a config is not found.
	_ = viper.ReadInConfig()

	// Set up the cobra root command
	var rootCmd = &cobra.Command{
		Use:   "sqlboiler [flags] <driver>",
		Short: "SQL Boiler generates an ORM tailored to your database schema.",
		Long: "SQL Boiler generates a Go ORM from template files, tailored to your database schema.\n" +
			`Complete documentation is available at http://github.com/vattle/sqlboiler`,
		Example:       `sqlboiler postgres`,
		PreRunE:       preRun,
		RunE:          run,
		PostRunE:      postRun,
		SilenceErrors: true,
		SilenceUsage:  true,
	}

	// Set up the cobra root command flags
	rootCmd.PersistentFlags().StringP("output", "o", "models", "The name of the folder to output to")
	rootCmd.PersistentFlags().StringP("schema", "s", "public", "The name of your database schema, for databases that support real schemas")
	rootCmd.PersistentFlags().StringP("pkgname", "p", "models", "The name you wish to assign to your generated package")
	rootCmd.PersistentFlags().StringP("basedir", "", "", "The base directory has the templates and templates_test folders")
	rootCmd.PersistentFlags().StringSliceP("blacklist", "b", nil, "Do not include these tables in your generated package")
	rootCmd.PersistentFlags().StringSliceP("whitelist", "w", nil, "Only include these tables in your generated package")
	rootCmd.PersistentFlags().StringSliceP("tag", "t", nil, "Struct tags to be included on your models in addition to json, yaml, toml")
	rootCmd.PersistentFlags().StringSliceP("replace", "", nil, "Replace templates by directory: relpath/to_file.tpl:relpath/to_replacement.tpl")
	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug mode prints stack traces on error")
	rootCmd.PersistentFlags().BoolP("no-tests", "", false, "Disable generated go test files")
	rootCmd.PersistentFlags().BoolP("no-hooks", "", false, "Disable hooks feature for your models")
	rootCmd.PersistentFlags().BoolP("no-auto-timestamps", "", false, "Disable automatic timestamps for created_at/updated_at")
	rootCmd.PersistentFlags().BoolP("version", "", false, "Print the version")
	rootCmd.PersistentFlags().BoolP("tinyint-as-bool", "", false, "Map MySQL tinyint(1) in Go to bool instead of int8")
	rootCmd.PersistentFlags().BoolP("wipe", "", false, "Delete the output folder (rm -rf) before generation to ensure sanity")

	viper.SetDefault("postgres.sslmode", "require")
	viper.SetDefault("postgres.port", "5432")
	viper.SetDefault("mysql.sslmode", "true")
	viper.SetDefault("mysql.port", "3306")

	viper.BindPFlags(rootCmd.PersistentFlags())
	viper.AutomaticEnv()

	if err := rootCmd.Execute(); err != nil {
		if e, ok := err.(commandFailure); ok {
			fmt.Printf("Error: %v\n\n", string(e))
			rootCmd.Help()
		} else if !viper.GetBool("debug") {
			fmt.Printf("Error: %v\n", err)
		} else {
			fmt.Printf("Error: %+v\n", err)
		}

		os.Exit(1)
	}
}

type commandFailure string

func (c commandFailure) Error() string {
	return string(c)
}

func preRun(cmd *cobra.Command, args []string) error {
	var err error

	if len(args) == 0 {
		return commandFailure("must provide a driver name")
	}

	driverName := args[0]

	cmdConfig = &boilingcore.Config{
		DriverName:       driverName,
		OutFolder:        viper.GetString("output"),
		Schema:           viper.GetString("schema"),
		PkgName:          viper.GetString("pkgname"),
		BaseDir:          viper.GetString("basedir"),
		Debug:            viper.GetBool("debug"),
		NoTests:          viper.GetBool("no-tests"),
		NoHooks:          viper.GetBool("no-hooks"),
		NoAutoTimestamps: viper.GetBool("no-auto-timestamps"),
		Wipe:             viper.GetBool("wipe"),
	}

	// BUG: https://github.com/spf13/viper/issues/200
	// Look up the value of blacklist, whitelist & tags directly from PFlags in Cobra if we
	// detect a malformed value coming out of viper.
	// Once the bug is fixed we'll be able to move this into the init above
	cmdConfig.BlacklistTables = viper.GetStringSlice("blacklist")
	if len(cmdConfig.BlacklistTables) == 1 && strings.ContainsRune(cmdConfig.BlacklistTables[0], ',') {
		cmdConfig.BlacklistTables, err = cmd.PersistentFlags().GetStringSlice("blacklist")
		if err != nil {
			return err
		}
	}

	cmdConfig.WhitelistTables = viper.GetStringSlice("whitelist")
	if len(cmdConfig.WhitelistTables) == 1 && strings.ContainsRune(cmdConfig.WhitelistTables[0], ',') {
		cmdConfig.WhitelistTables, err = cmd.PersistentFlags().GetStringSlice("whitelist")
		if err != nil {
			return err
		}
	}

	cmdConfig.Tags = viper.GetStringSlice("tag")
	if len(cmdConfig.Tags) == 1 && strings.ContainsRune(cmdConfig.Tags[0], ',') {
		cmdConfig.Tags, err = cmd.PersistentFlags().GetStringSlice("tag")
		if err != nil {
			return err
		}
	}

	cmdConfig.Replacements = viper.GetStringSlice("replace")
	if len(cmdConfig.Replacements) == 1 && strings.ContainsRune(cmdConfig.Replacements[0], ',') {
		cmdConfig.Replacements, err = cmd.PersistentFlags().GetStringSlice("replace")
		if err != nil {
			return err
		}
	}

	if driverName == "postgres" {
		cmdConfig.Postgres = boilingcore.PostgresConfig{
			User:    viper.GetString("postgres.user"),
			Pass:    viper.GetString("postgres.pass"),
			Host:    viper.GetString("postgres.host"),
			Port:    viper.GetInt("postgres.port"),
			DBName:  viper.GetString("postgres.dbname"),
			SSLMode: viper.GetString("postgres.sslmode"),
		}

		// BUG: https://github.com/spf13/viper/issues/71
		// Despite setting defaults, nested values don't get defaults
		// Set them manually
		if cmdConfig.Postgres.SSLMode == "" {
			cmdConfig.Postgres.SSLMode = "require"
			viper.Set("postgres.sslmode", cmdConfig.Postgres.SSLMode)
		}

		if cmdConfig.Postgres.Port == 0 {
			cmdConfig.Postgres.Port = 5432
			viper.Set("postgres.port", cmdConfig.Postgres.Port)
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.Postgres.User, "postgres.user"),
			vala.StringNotEmpty(cmdConfig.Postgres.Host, "postgres.host"),
			vala.Not(vala.Equals(cmdConfig.Postgres.Port, 0, "postgres.port")),
			vala.StringNotEmpty(cmdConfig.Postgres.DBName, "postgres.dbname"),
			vala.StringNotEmpty(cmdConfig.Postgres.SSLMode, "postgres.sslmode"),
		).Check()

		if err != nil {
			return commandFailure(err.Error())
		}
	}

	if driverName == "mysql" {
		cmdConfig.MySQL = boilingcore.MySQLConfig{
			User:    viper.GetString("mysql.user"),
			Pass:    viper.GetString("mysql.pass"),
			Host:    viper.GetString("mysql.host"),
			Port:    viper.GetInt("mysql.port"),
			DBName:  viper.GetString("mysql.dbname"),
			SSLMode: viper.GetString("mysql.sslmode"),
		}

		// Set MySQL TinyintAsBool global var. This flag only applies to MySQL.
		drivers.TinyintAsBool = viper.GetBool("tinyint-as-bool")

		// MySQL doesn't have schemas, just databases
		cmdConfig.Schema = cmdConfig.MySQL.DBName

		// BUG: https://github.com/spf13/viper/issues/71
		// Despite setting defaults, nested values don't get defaults
		// Set them manually
		if cmdConfig.MySQL.SSLMode == "" {
			cmdConfig.MySQL.SSLMode = "true"
			viper.Set("mysql.sslmode", cmdConfig.MySQL.SSLMode)
		}

		if cmdConfig.MySQL.Port == 0 {
			cmdConfig.MySQL.Port = 3306
			viper.Set("mysql.port", cmdConfig.MySQL.Port)
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.MySQL.User, "mysql.user"),
			vala.StringNotEmpty(cmdConfig.MySQL.Host, "mysql.host"),
			vala.Not(vala.Equals(cmdConfig.MySQL.Port, 0, "mysql.port")),
			vala.StringNotEmpty(cmdConfig.MySQL.DBName, "mysql.dbname"),
			vala.StringNotEmpty(cmdConfig.MySQL.SSLMode, "mysql.sslmode"),
		).Check()

		if err != nil {
			return commandFailure(err.Error())
		}
	}

	cmdState, err = boilingcore.New(cmdConfig)
	return err
}

func run(cmd *cobra.Command, args []string) error {
	return cmdState.Run(true)
}

func postRun(cmd *cobra.Command, args []string) error {
	return cmdState.Cleanup()
}
<commit_msg>Hide replace flag from help output<commit_after>// Package main defines a command line interface for the sqlboiler package
package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/kat-co/vala"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/vattle/sqlboiler/bdb/drivers"
	"github.com/vattle/sqlboiler/boilingcore"
)

const sqlBoilerVersion = "2.1.9"

var (
	cmdState  *boilingcore.State
	cmdConfig *boilingcore.Config
)

func main() {
	var err error

	// Too much happens between here and cobra's argument handling, for
	// something so simple just do it immediately.
	for _, arg := range os.Args {
		if arg == "--version" {
			fmt.Println("SQLBoiler v" + sqlBoilerVersion)
			return
		}
	}

	viper.SetConfigName("sqlboiler")

	configHome := os.Getenv("XDG_CONFIG_HOME")
	homePath := os.Getenv("HOME")
	wd, err := os.Getwd()
	if err != nil {
		wd = "./"
	}

	configPaths := []string{wd}
	if len(configHome) > 0 {
		configPaths = append(configPaths, filepath.Join(configHome, "sqlboiler"))
	} else {
		configPaths = append(configPaths, filepath.Join(homePath, ".config/sqlboiler"))
	}

	for _, p := range configPaths {
		viper.AddConfigPath(p)
	}

	// Ignore errors here, fallback to other validation methods.
	// Users can use environment variables if a config is not found.
	_ = viper.ReadInConfig()

	// Set up the cobra root command
	var rootCmd = &cobra.Command{
		Use:   "sqlboiler [flags] <driver>",
		Short: "SQL Boiler generates an ORM tailored to your database schema.",
		Long: "SQL Boiler generates a Go ORM from template files, tailored to your database schema.\n" +
			`Complete documentation is available at http://github.com/vattle/sqlboiler`,
		Example:       `sqlboiler postgres`,
		PreRunE:       preRun,
		RunE:          run,
		PostRunE:      postRun,
		SilenceErrors: true,
		SilenceUsage:  true,
	}

	// Set up the cobra root command flags
	rootCmd.PersistentFlags().StringP("output", "o", "models", "The name of the folder to output to")
	rootCmd.PersistentFlags().StringP("schema", "s", "public", "The name of your database schema, for databases that support real schemas")
	rootCmd.PersistentFlags().StringP("pkgname", "p", "models", "The name you wish to assign to your generated package")
	rootCmd.PersistentFlags().StringP("basedir", "", "", "The base directory has the templates and templates_test folders")
	rootCmd.PersistentFlags().StringSliceP("blacklist", "b", nil, "Do not include these tables in your generated package")
	rootCmd.PersistentFlags().StringSliceP("whitelist", "w", nil, "Only include these tables in your generated package")
	rootCmd.PersistentFlags().StringSliceP("tag", "t", nil, "Struct tags to be included on your models in addition to json, yaml, toml")
	rootCmd.PersistentFlags().StringSliceP("replace", "", nil, "Replace templates by directory: relpath/to_file.tpl:relpath/to_replacement.tpl")
	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug mode prints stack traces on error")
	rootCmd.PersistentFlags().BoolP("no-tests", "", false, "Disable generated go test files")
	rootCmd.PersistentFlags().BoolP("no-hooks", "", false, "Disable hooks feature for your models")
	rootCmd.PersistentFlags().BoolP("no-auto-timestamps", "", false, "Disable automatic timestamps for created_at/updated_at")
	rootCmd.PersistentFlags().BoolP("version", "", false, "Print the version")
	rootCmd.PersistentFlags().BoolP("tinyint-as-bool", "", false, "Map MySQL tinyint(1) in Go to bool instead of int8")
	rootCmd.PersistentFlags().BoolP("wipe", "", false, "Delete the output folder (rm -rf) before generation to ensure sanity")

	// hide flags not recommended for use
	rootCmd.PersistentFlags().MarkHidden("replace")

	viper.SetDefault("postgres.sslmode", "require")
	viper.SetDefault("postgres.port", "5432")
	viper.SetDefault("mysql.sslmode", "true")
	viper.SetDefault("mysql.port", "3306")

	viper.BindPFlags(rootCmd.PersistentFlags())
	viper.AutomaticEnv()

	if err := rootCmd.Execute(); err != nil {
		if e, ok := err.(commandFailure); ok {
			fmt.Printf("Error: %v\n\n", string(e))
			rootCmd.Help()
		} else if !viper.GetBool("debug") {
			fmt.Printf("Error: %v\n", err)
		} else {
			fmt.Printf("Error: %+v\n", err)
		}

		os.Exit(1)
	}
}

type commandFailure string

func (c commandFailure) Error() string {
	return string(c)
}

func preRun(cmd *cobra.Command, args []string) error {
	var err error

	if len(args) == 0 {
		return commandFailure("must provide a driver name")
	}

	driverName := args[0]

	cmdConfig = &boilingcore.Config{
		DriverName:       driverName,
		OutFolder:        viper.GetString("output"),
		Schema:           viper.GetString("schema"),
		PkgName:          viper.GetString("pkgname"),
		BaseDir:          viper.GetString("basedir"),
		Debug:            viper.GetBool("debug"),
		NoTests:          viper.GetBool("no-tests"),
		NoHooks:          viper.GetBool("no-hooks"),
		NoAutoTimestamps: viper.GetBool("no-auto-timestamps"),
		Wipe:             viper.GetBool("wipe"),
	}

	// BUG: https://github.com/spf13/viper/issues/200
	// Look up the value of blacklist, whitelist & tags directly from PFlags in Cobra if we
	// detect a malformed value coming out of viper.
	// Once the bug is fixed we'll be able to move this into the init above
	cmdConfig.BlacklistTables = viper.GetStringSlice("blacklist")
	if len(cmdConfig.BlacklistTables) == 1 && strings.ContainsRune(cmdConfig.BlacklistTables[0], ',') {
		cmdConfig.BlacklistTables, err = cmd.PersistentFlags().GetStringSlice("blacklist")
		if err != nil {
			return err
		}
	}

	cmdConfig.WhitelistTables = viper.GetStringSlice("whitelist")
	if len(cmdConfig.WhitelistTables) == 1 && strings.ContainsRune(cmdConfig.WhitelistTables[0], ',') {
		cmdConfig.WhitelistTables, err = cmd.PersistentFlags().GetStringSlice("whitelist")
		if err != nil {
			return err
		}
	}

	cmdConfig.Tags = viper.GetStringSlice("tag")
	if len(cmdConfig.Tags) == 1 && strings.ContainsRune(cmdConfig.Tags[0], ',') {
		cmdConfig.Tags, err = cmd.PersistentFlags().GetStringSlice("tag")
		if err != nil {
			return err
		}
	}

	cmdConfig.Replacements = viper.GetStringSlice("replace")
	if len(cmdConfig.Replacements) == 1 && strings.ContainsRune(cmdConfig.Replacements[0], ',') {
		cmdConfig.Replacements, err = cmd.PersistentFlags().GetStringSlice("replace")
		if err != nil {
			return err
		}
	}

	if driverName == "postgres" {
		cmdConfig.Postgres = boilingcore.PostgresConfig{
			User:    viper.GetString("postgres.user"),
			Pass:    viper.GetString("postgres.pass"),
			Host:    viper.GetString("postgres.host"),
			Port:    viper.GetInt("postgres.port"),
			DBName:  viper.GetString("postgres.dbname"),
			SSLMode: viper.GetString("postgres.sslmode"),
		}

		// BUG: https://github.com/spf13/viper/issues/71
		// Despite setting defaults, nested values don't get defaults
		// Set them manually
		if cmdConfig.Postgres.SSLMode == "" {
			cmdConfig.Postgres.SSLMode = "require"
			viper.Set("postgres.sslmode", cmdConfig.Postgres.SSLMode)
		}

		if cmdConfig.Postgres.Port == 0 {
			cmdConfig.Postgres.Port = 5432
			viper.Set("postgres.port", cmdConfig.Postgres.Port)
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.Postgres.User, "postgres.user"),
			vala.StringNotEmpty(cmdConfig.Postgres.Host, "postgres.host"),
			vala.Not(vala.Equals(cmdConfig.Postgres.Port, 0, "postgres.port")),
			vala.StringNotEmpty(cmdConfig.Postgres.DBName, "postgres.dbname"),
			vala.StringNotEmpty(cmdConfig.Postgres.SSLMode, "postgres.sslmode"),
		).Check()

		if err != nil {
			return commandFailure(err.Error())
		}
	}

	if driverName == "mysql" {
		cmdConfig.MySQL = boilingcore.MySQLConfig{
			User:    viper.GetString("mysql.user"),
			Pass:    viper.GetString("mysql.pass"),
			Host:    viper.GetString("mysql.host"),
			Port:    viper.GetInt("mysql.port"),
			DBName:  viper.GetString("mysql.dbname"),
			SSLMode: viper.GetString("mysql.sslmode"),
		}

		// Set MySQL TinyintAsBool global var. This flag only applies to MySQL.
		drivers.TinyintAsBool = viper.GetBool("tinyint-as-bool")

		// MySQL doesn't have schemas, just databases
		cmdConfig.Schema = cmdConfig.MySQL.DBName

		// BUG: https://github.com/spf13/viper/issues/71
		// Despite setting defaults, nested values don't get defaults
		// Set them manually
		if cmdConfig.MySQL.SSLMode == "" {
			cmdConfig.MySQL.SSLMode = "true"
			viper.Set("mysql.sslmode", cmdConfig.MySQL.SSLMode)
		}

		if cmdConfig.MySQL.Port == 0 {
			cmdConfig.MySQL.Port = 3306
			viper.Set("mysql.port", cmdConfig.MySQL.Port)
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.MySQL.User, "mysql.user"),
			vala.StringNotEmpty(cmdConfig.MySQL.Host, "mysql.host"),
			vala.Not(vala.Equals(cmdConfig.MySQL.Port, 0, "mysql.port")),
			vala.StringNotEmpty(cmdConfig.MySQL.DBName, "mysql.dbname"),
			vala.StringNotEmpty(cmdConfig.MySQL.SSLMode, "mysql.sslmode"),
		).Check()

		if err != nil {
			return commandFailure(err.Error())
		}
	}

	cmdState, err = boilingcore.New(cmdConfig)
	return err
}

func run(cmd *cobra.Command, args []string) error {
	return cmdState.Run(true)
}

func postRun(cmd *cobra.Command, args []string) error {
	return cmdState.Cleanup()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"syscall"
	"unsafe"
)

var (
	modshell32         = syscall.NewLazyDLL("shell32.dll")
	procShellExecuteEx = modshell32.NewProc("ShellExecuteExW")
)

const (
	SEE_MASK_NOCLOSEPROCESS = 0x00000040
)

const (
	ERROR_BAD_FORMAT = 11
)

const (
	SE_ERR_FNF             = 2
	SE_ERR_PNF             = 3
	SE_ERR_ACCESSDENIED    = 5
	SE_ERR_OOM             = 8
	SE_ERR_DLLNOTFOUND     = 32
	SE_ERR_SHARE           = 26
	SE_ERR_ASSOCINCOMPLETE = 27
	SE_ERR_DDETIMEOUT      = 28
	SE_ERR_DDEFAIL         = 29
	SE_ERR_DDEBUSY         = 30
	SE_ERR_NOASSOC         = 31
)

type (
	DWORD     uint32
	HANDLE    uintptr
	HINSTANCE HANDLE
	HKEY      HANDLE
	HWND      HANDLE
	ULONG     uint32
	LPCTSTR   uintptr
	LPVOID    uintptr
)

type SHELLEXECUTEINFO struct {
	cbSize         DWORD
	fMask          ULONG
	hwnd           HWND
	lpVerb         LPCTSTR
	lpFile         LPCTSTR
	lpParameters   LPCTSTR
	lpDirectory    LPCTSTR
	nShow          int
	hInstApp       HINSTANCE
	lpIDList       LPVOID
	lpClass        LPCTSTR
	hkeyClass      HKEY
	dwHotKey       DWORD
	hIconOrMonitor HANDLE
	hProcess       HANDLE
}

func ShellExecuteAndWait(hwnd HWND, lpOperation, lpFile, lpParameters, lpDirectory string, nShowCmd int) error {
	var lpctstrVerb, lpctstrParameters, lpctstrDirectory LPCTSTR
	if len(lpOperation) != 0 {
		lpctstrVerb = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpOperation)))
	}
	if len(lpParameters) != 0 {
		lpctstrParameters = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpParameters)))
	}
	if len(lpDirectory) != 0 {
		lpctstrDirectory = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDirectory)))
	}
	i := &SHELLEXECUTEINFO{
		fMask:        SEE_MASK_NOCLOSEPROCESS,
		hwnd:         hwnd,
		lpVerb:       lpctstrVerb,
		lpFile:       LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpFile))),
		lpParameters: lpctstrParameters,
		lpDirectory:  lpctstrDirectory,
		nShow:        nShowCmd,
	}
	i.cbSize = DWORD(unsafe.Sizeof(*i))
	return ShellExecuteEx(i)
}

func ShellExecuteEx(pExecInfo *SHELLEXECUTEINFO) error {
	ret, _, _ := procShellExecuteEx.Call(uintptr(unsafe.Pointer(pExecInfo)))
	if ret == 1 && pExecInfo.fMask&SEE_MASK_NOCLOSEPROCESS != 0 {
		s, e := syscall.WaitForSingleObject(syscall.Handle(pExecInfo.hProcess), syscall.INFINITE)
		switch s {
		case syscall.WAIT_OBJECT_0:
			break
		case syscall.WAIT_FAILED:
			return os.NewSyscallError("WaitForSingleObject", e)
		default:
			return errors.New("Unexpected result from WaitForSingleObject")
		}
	}
	errorMsg := ""
	if pExecInfo.hInstApp != 0 && pExecInfo.hInstApp <= 32 {
		switch int(pExecInfo.hInstApp) {
		case SE_ERR_FNF:
			errorMsg = "The specified file was not found"
		case SE_ERR_PNF:
			errorMsg = "The specified path was not found"
		case ERROR_BAD_FORMAT:
			errorMsg = "The .exe file is invalid (non-Win32 .exe or error in .exe image)"
		case SE_ERR_ACCESSDENIED:
			errorMsg = "The operating system denied access to the specified file"
		case SE_ERR_ASSOCINCOMPLETE:
			errorMsg = "The file name association is incomplete or invalid"
		case SE_ERR_DDEBUSY:
			errorMsg = "The DDE transaction could not be completed because other DDE transactions were being processed"
		case SE_ERR_DDEFAIL:
			errorMsg = "The DDE transaction failed"
		case SE_ERR_DDETIMEOUT:
			errorMsg = "The DDE transaction could not be completed because the request timed out"
		case SE_ERR_DLLNOTFOUND:
			errorMsg = "The specified DLL was not found"
		case SE_ERR_NOASSOC:
			errorMsg = "There is no application associated with the given file name extension"
		case SE_ERR_OOM:
			errorMsg = "There was not enough memory to complete the operation"
		case SE_ERR_SHARE:
			errorMsg = "A sharing violation occurred"
		default:
			errorMsg = fmt.Sprintf("Unknown error occurred with error code %v", pExecInfo.hInstApp)
		}
	} else {
		return nil
	}
	return errors.New(errorMsg)
}

type msg struct {
	Name string
	Exit int
	Data []byte
}

func msgWrite(enc *gob.Encoder, typ string) io.WriteCloser {
	r, w := io.Pipe()
	go func() {
		defer r.Close()
		var b [4096]byte
		for {
			n, err := r.Read(b[:])
			if err != nil {
				break
			}
			enc.Encode(&msg{Name: typ, Data: b[:n]})
		}
	}()
	return w
}

func client(addr string) int {
	// connect to server
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer conn.Close()
	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)

	// stdin
	inw, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer inw.Close()

	// stdout
	outw := msgWrite(enc, "stdout")
	defer outw.Close()
	cmd.Stdout = outw

	// stderr
	errw := msgWrite(enc, "stderr")
	defer errw.Close()
	cmd.Stderr = errw

	go func() {
		defer inw.Close()
	in_loop:
		for {
			var m msg
			err = dec.Decode(&m)
			if err != nil {
				return
			}
			switch m.Name {
			case "close":
				break in_loop
			case "ctrlc":
				if runtime.GOOS == "windows" {
					// windows doesn't support os.Interrupt
					cmd.Process.Kill()
				} else {
					cmd.Process.Signal(os.Interrupt)
				}
			case "stdin":
				inw.Write(m.Data)
			}
		}
	}()

	err = cmd.Run()

	code := 1
	if err != nil {
		if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
			code = status.ExitStatus()
		}
	} else {
		code = 0
	}

	enc.Encode(&msg{Name: "exit", Exit: code})
	return 0
}

func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += syscall.EscapeArg(v)
	}
	return s
}

func server() int {
	// make listner to communicate child process
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer lis.Close()

	// make sure executable name to avoid detecting same executable name
	exe, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	args := []string{"-mode", lis.Addr().String()}
	args = append(args, flag.Args()...)

	go func() {
		err = ShellExecuteAndWait(0, "runas", exe, makeCmdLine(args), "", syscall.SW_HIDE)
		if err != nil {
			lis.Close()
		}
	}()

	conn, err := lis.Accept()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
		return 1
	}
	defer conn.Close()

	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, os.Interrupt)
	go func() {
		for range sc {
			enc.Encode(&msg{Name: "ctrlc"})
		}
	}()
	defer close(sc)

	go func() {
		var b [256]byte
		for {
			n, err := os.Stdin.Read(b[:])
			if err != nil {
				// stdin was closed
				enc.Encode(&msg{Name: "close"})
				break
			}
			enc.Encode(&msg{Name: "stdin", Data: b[:n]})
		}
	}()

	for {
		var m msg
		err = dec.Decode(&m)
		if err != nil {
			fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
			return 1
		}
		switch m.Name {
		case "stdout":
			os.Stdout.Write(m.Data)
		case "stderr":
			os.Stderr.Write(m.Data)
		case "exit":
			return m.Exit
		}
	}
	return 0
}

func main() {
	var mode string
	flag.StringVar(&mode, "mode", "", "mode")
	flag.Parse()
	if mode != "" {
		os.Exit(client(mode))
	}
	os.Exit(server())
}
<commit_msg>use syscall because go handle unicode strings<commit_after>package main

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"syscall"
	"unsafe"
)

var (
	modshell32         = syscall.NewLazyDLL("shell32.dll")
	procShellExecuteEx = modshell32.NewProc("ShellExecuteExW")
)

const (
	SEE_MASK_NOCLOSEPROCESS = 0x00000040
)

const (
	ERROR_BAD_FORMAT = 11
)

const (
	SE_ERR_FNF             = 2
	SE_ERR_PNF             = 3
	SE_ERR_ACCESSDENIED    = 5
	SE_ERR_OOM             = 8
	SE_ERR_DLLNOTFOUND     = 32
	SE_ERR_SHARE           = 26
	SE_ERR_ASSOCINCOMPLETE = 27
	SE_ERR_DDETIMEOUT      = 28
	SE_ERR_DDEFAIL         = 29
	SE_ERR_DDEBUSY         = 30
	SE_ERR_NOASSOC         = 31
)

type (
	DWORD     uint32
	HANDLE    uintptr
	HINSTANCE HANDLE
	HKEY      HANDLE
	HWND      HANDLE
	ULONG     uint32
	LPCTSTR   uintptr
	LPVOID    uintptr
)

type SHELLEXECUTEINFO struct {
	cbSize         DWORD
	fMask          ULONG
	hwnd           HWND
	lpVerb         LPCTSTR
	lpFile         LPCTSTR
	lpParameters   LPCTSTR
	lpDirectory    LPCTSTR
	nShow          int
	hInstApp       HINSTANCE
	lpIDList       LPVOID
	lpClass        LPCTSTR
	hkeyClass      HKEY
	dwHotKey       DWORD
	hIconOrMonitor HANDLE
	hProcess       HANDLE
}

func ShellExecuteAndWait(hwnd HWND, lpOperation, lpFile, lpParameters, lpDirectory string, nShowCmd int) error {
	var lpctstrVerb, lpctstrParameters, lpctstrDirectory LPCTSTR
	if len(lpOperation) != 0 {
		lpctstrVerb = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpOperation)))
	}
	if len(lpParameters) != 0 {
		lpctstrParameters = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpParameters)))
	}
	if len(lpDirectory) != 0 {
		lpctstrDirectory = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDirectory)))
	}
	i := &SHELLEXECUTEINFO{
		fMask:        SEE_MASK_NOCLOSEPROCESS,
		hwnd:         hwnd,
		lpVerb:       lpctstrVerb,
		lpFile:       LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpFile))),
		lpParameters: lpctstrParameters,
		lpDirectory:  lpctstrDirectory,
		nShow:        nShowCmd,
	}
	i.cbSize = DWORD(unsafe.Sizeof(*i))
	return ShellExecuteEx(i)
}

func ShellExecuteEx(pExecInfo *SHELLEXECUTEINFO) error {
	ret, _, _ := procShellExecuteEx.Call(uintptr(unsafe.Pointer(pExecInfo)))
	if ret == 1 && pExecInfo.fMask&SEE_MASK_NOCLOSEPROCESS != 0 {
		s, e := syscall.WaitForSingleObject(syscall.Handle(pExecInfo.hProcess), syscall.INFINITE)
		switch s {
		case syscall.WAIT_OBJECT_0:
			break
		case syscall.WAIT_FAILED:
			return os.NewSyscallError("WaitForSingleObject", e)
		default:
			return errors.New("Unexpected result from WaitForSingleObject")
		}
	}
	errorMsg := ""
	if pExecInfo.hInstApp != 0 && pExecInfo.hInstApp <= 32 {
		switch int(pExecInfo.hInstApp) {
		case SE_ERR_FNF:
			errorMsg = "The specified file was not found"
		case SE_ERR_PNF:
			errorMsg = "The specified path was not found"
		case ERROR_BAD_FORMAT:
			errorMsg = "The .exe file is invalid (non-Win32 .exe or error in .exe image)"
		case SE_ERR_ACCESSDENIED:
			errorMsg = "The operating system denied access to the specified file"
		case SE_ERR_ASSOCINCOMPLETE:
			errorMsg = "The file name association is incomplete or invalid"
		case SE_ERR_DDEBUSY:
			errorMsg = "The DDE transaction could not be completed because other DDE transactions were being processed"
		case SE_ERR_DDEFAIL:
			errorMsg = "The DDE transaction failed"
		case SE_ERR_DDETIMEOUT:
			errorMsg = "The DDE transaction could not be completed because the request timed out"
		case SE_ERR_DLLNOTFOUND:
			errorMsg = "The specified DLL was not found"
		case SE_ERR_NOASSOC:
			errorMsg = "There is no application associated with the given file name extension"
		case SE_ERR_OOM:
			errorMsg = "There was not enough memory to complete the operation"
		case SE_ERR_SHARE:
			errorMsg = "A sharing violation occurred"
		default:
			errorMsg = fmt.Sprintf("Unknown error occurred with error code %v", pExecInfo.hInstApp)
		}
	} else {
		return nil
	}
	return errors.New(errorMsg)
}

type msg struct {
	Name string
	Exit int
	Data []byte
}

func msgWrite(enc *gob.Encoder, typ string) io.WriteCloser {
	r, w := io.Pipe()
	go func() {
		defer r.Close()
		var b [4096]byte
		for {
			n, err := r.Read(b[:])
			if err != nil {
				break
			}
			enc.Encode(&msg{Name: typ, Data: b[:n]})
		}
	}()
	return w
}

func client(addr string) int {
	// connect to server
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer conn.Close()
	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)

	// stdin
	inw, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer inw.Close()

	// stdout
	outw := msgWrite(enc, "stdout")
	defer outw.Close()
	cmd.Stdout = outw

	// stderr
	errw := msgWrite(enc, "stderr")
	defer errw.Close()
	cmd.Stderr = errw

	go func() {
		defer inw.Close()
	in_loop:
		for {
			var m msg
			err = dec.Decode(&m)
			if err != nil {
				return
			}
			switch m.Name {
			case "close":
				break in_loop
			case "ctrlc":
				if runtime.GOOS == "windows" {
					// windows doesn't support os.Interrupt
					cmd.Process.Kill()
				} else {
					cmd.Process.Signal(os.Interrupt)
				}
			case "stdin":
				inw.Write(m.Data)
			}
		}
	}()

	err = cmd.Run()

	code := 1
	if err != nil {
		if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
			code = status.ExitStatus()
		}
	} else {
		code = 0
	}

	enc.Encode(&msg{Name: "exit", Exit: code})
	return 0
}

func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += syscall.EscapeArg(v)
	}
	return s
}

func server() int {
	// make listner to communicate child process
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer lis.Close()

	// make sure executable name to avoid detecting same executable name
	exe, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	args := []string{"-mode", lis.Addr().String()}
	args = append(args, flag.Args()...)

	go func() {
		err = ShellExecuteAndWait(0, "runas", exe, makeCmdLine(args), "", syscall.SW_HIDE)
		if err != nil {
			lis.Close()
		}
	}()

	conn, err := lis.Accept()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
		return 1
	}
	defer conn.Close()

	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, os.Interrupt)
	go func() {
		for range sc {
			enc.Encode(&msg{Name: "ctrlc"})
		}
	}()
	defer close(sc)

	go func() {
		var b [256]byte
		for {
			n, err := os.Stdin.Read(b[:])
			if err != nil {
				// stdin was closed
				enc.Encode(&msg{Name: "close"})
				break
			}
			enc.Encode(&msg{Name: "stdin", Data: b[:n]})
		}
	}()

	for {
		var m msg
		err = dec.Decode(&m)
		if err != nil {
			fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
			return 1
		}
		switch m.Name {
		case "stdout":
			syscall.Write(syscall.Stdout, m.Data)
		case "stderr":
			syscall.Write(syscall.Stderr, m.Data)
		case "exit":
			return m.Exit
		}
	}
	return 0
}

func main() {
	var mode string
	flag.StringVar(&mode, "mode", "", "mode")
	flag.Parse()
	if mode != "" {
		os.Exit(client(mode))
	}
	os.Exit(server())
}
<|endoftext|>


<commit_before>// Copyright 2016 Tamás Gulácsi. All rights reserved.
// Use of this source code is governed by The MIT License
// found in the accompanying LICENSE file.

// Package date implements encoding of 7 byte Oracle DATE storage formats.
package date

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"
)

// Date is an OCIDate
//
// SQLT_ODT: 7 bytes
//
// http://www.orafaq.com/wiki/Date
//
/*
The internal format is the following one:

    century + 100
    year in the century + 100
    month
    day
    hour + 1
    minute + 1
    second + 1

So in the previous example the date was 19-DEC-2007 at 22:35:10.
*/
type Date [7]byte

func (dt *Date) Set(t time.Time) {
	if t.IsZero() {
		for i := range dt[:] {
			dt[i] = 0
		}
		return
	}
	y := t.Year()
	if y < -4711 {
		y = -4711
	} else if y > 9999 {
		y = 9999
	}
	dt[0] = byte(y/100 + 100)
	dt[1] = byte(y%100 + 100)
	dt[2] = byte(t.Month())
	dt[3] = byte(t.Day())
	dt[4] = byte(t.Hour() + 1)
	dt[5] = byte(t.Minute() + 1)
	dt[6] = byte(t.Second() + 1)
}

func (dt *Date) IsNull() bool {
	for _, b := range dt[:] {
		if b != 0 {
			return false
		}
	}
	return true
}
func (dt Date) MarshalJSON() ([]byte, error) {
	if dt.IsNull() {
		return []byte("null"), nil
	}
	return dt.Get().MarshalJSON()
}
func (dt *Date) UnmarshalJSON(p []byte) error {
	if bytes.Equal(p, []byte("null")) || bytes.Equal(p, []byte(`""`)) {
		for i := range dt[:] {
			dt[i] = 0
		}
		return nil
	}
	var t time.Time
	if err := json.Unmarshal(p, &t); err != nil {
		return err
	}
	dt.Set(t)
	return nil
}

// FromTime returns a Date from a time.Time
// Does the allocation inside, so easier to use.
func FromTime(t time.Time) Date {
	var dt Date
	dt.Set(t)
	return dt
}

func (dt Date) Equal(other Date) bool {
	return bytes.Equal(dt[:], other[:])
}

func (dt Date) String() string {
	if dt.IsNull() {
		return (time.Time{}).Format("2006-01-02T15:04:05")
	}
	return fmt.Sprintf("%04d-%02d-%02dT%02d:%02d:%02d",
		(int(dt[0])-100)*100+(int(dt[1])-100),
		time.Month(dt[2]),
		int(dt[3]),
		int(dt[4]-1),
		int(dt[5]-1),
		int(dt[6]-1),
	)
}

func (dt Date) Get() time.Time {
	return dt.GetIn(nil)
}
func (dt Date) GetIn(zone *time.Location) time.Time {
	if dt.IsNull() {
		return time.Time{}
	}
	if zone == nil {
		zone = time.Local
	}
	return time.Date(
		(int(dt[0])-100)*100+(int(dt[1])-100),
		time.Month(dt[2]),
		int(dt[3]),
		int(dt[4]-1),
		int(dt[5]-1),
		int(dt[6]-1),
		0,
		zone,
	)
}
<commit_msg>Date: add Bytes() function<commit_after>// Copyright 2016 Tamás Gulácsi. All rights reserved.
// Use of this source code is governed by The MIT License
// found in the accompanying LICENSE file.

// Package date implements encoding of 7 byte Oracle DATE storage formats.
package date

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"
)

// Date is an OCIDate
//
// SQLT_ODT: 7 bytes
//
// http://www.orafaq.com/wiki/Date
//
/*
The internal format is the following one:

    century + 100
    year in the century + 100
    month
    day
    hour + 1
    minute + 1
    second + 1

So in the previous example the date was 19-DEC-2007 at 22:35:10.
*/
type Date [7]byte

func (dt *Date) Set(t time.Time) {
	if t.IsZero() {
		for i := range dt[:] {
			dt[i] = 0
		}
		return
	}
	y := t.Year()
	if y < -4711 {
		y = -4711
	} else if y > 9999 {
		y = 9999
	}
	dt[0] = byte(y/100 + 100)
	dt[1] = byte(y%100 + 100)
	dt[2] = byte(t.Month())
	dt[3] = byte(t.Day())
	dt[4] = byte(t.Hour() + 1)
	dt[5] = byte(t.Minute() + 1)
	dt[6] = byte(t.Second() + 1)
}

func (dt Date) Bytes() []byte {
	return dt[:]
}

func (dt Date) IsNull() bool {
	for _, b := range dt[:] {
		if b != 0 {
			return false
		}
	}
	return true
}
func (dt Date) MarshalJSON() ([]byte, error) {
	if dt.IsNull() {
		return []byte("null"), nil
	}
	return dt.Get().MarshalJSON()
}
func (dt *Date) UnmarshalJSON(p []byte) error {
	if bytes.Equal(p, []byte("null")) || bytes.Equal(p, []byte(`""`)) {
		for i := range dt[:] {
			dt[i] = 0
		}
		return nil
	}
	var t time.Time
	if err := json.Unmarshal(p, &t); err != nil {
		return err
	}
	dt.Set(t)
	return nil
}

// FromTime returns a Date from a time.Time
// Does the allocation inside, so easier to use.
func FromTime(t time.Time) Date {
	var dt Date
	dt.Set(t)
	return dt
}

func (dt Date) Equal(other Date) bool {
	return bytes.Equal(dt[:], other[:])
}

func (dt Date) String() string {
	if dt.IsNull() {
		return (time.Time{}).Format("2006-01-02T15:04:05")
	}
	return fmt.Sprintf("%04d-%02d-%02dT%02d:%02d:%02d",
		(int(dt[0])-100)*100+(int(dt[1])-100),
		time.Month(dt[2]),
		int(dt[3]),
		int(dt[4]-1),
		int(dt[5]-1),
		int(dt[6]-1),
	)
}

func (dt Date) Get() time.Time {
	return dt.GetIn(nil)
}
func (dt Date) GetIn(zone *time.Location) time.Time {
	if dt.IsNull() {
		return time.Time{}
	}
	if zone == nil {
		zone = time.Local
	}
	return time.Date(
		(int(dt[0])-100)*100+(int(dt[1])-100),
		time.Month(dt[2]),
		int(dt[3]),
		int(dt[4]-1),
		int(dt[5]-1),
		int(dt[6]-1),
		0,
		zone,
	)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/textproto"
	"os"
	"strconv"

	"github.com/Omar-H/goimage/util"
	"github.com/gorilla/csrf"
	"github.com/gorilla/mux"
)

// Data parsed to the ViewHandler's template
type ViewData struct {
	Id       string
	ImageUrl string
	Ext      string
}

// Parsed JSON config
type Config struct {
	Port              int      `json:"port"`
	Secure            bool     `json:"secure"`
	AuthKey           string   `json:"32-byte-auth-key"`
	AllowedMimeTypes  []string `json:"allowed-mime-types"`
	AllowedExtensions []string `json:"allowed-extensions"`
	ImageNameLength   int      `json:"image-name-length"`
	MaxFileSize       int64    `json:"max-file-size"`
	ImageDirectory    string   `json:"image-directory"`
	TemplateDirectory string   `json:"template-directory"`
	PublicDirectory   string   `json:"public-directory"`
	ImageUrl          string   `json:"image-url"`
	CSRF              bool     `json:"csrf"`
}

var config Config

func main() {
	jsonFiles, err := ioutil.ReadFile("./config.json")
	if err != nil {
		log.Fatal(err)
		return
	}
	json.Unmarshal(jsonFiles, &config)

	CSRF := csrf.Protect(
		[]byte(config.AuthKey),
		csrf.RequestHeader("X-CSRF-Token"),
		csrf.FieldName("_csrf"),
		csrf.Secure(config.Secure),
	)
	r := mux.NewRouter()
	err = os.MkdirAll(config.TemplateDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Template Directory")
	}
	err = os.MkdirAll(config.PublicDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Public Directory")
	}
	err = os.MkdirAll(config.ImageDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Image Directory")
	}

	templates := template.Must(template.ParseGlob(config.TemplateDirectory + "*.html"))

	r.HandleFunc("/{id}/", ViewHandler(templates)).Methods("GET")
	r.HandleFunc("/upload/", UploadHandler).Methods("POST")
	r.PathPrefix("/").HandlerFunc(RootHandler(templates)).Methods("GET")

	log.Print("Listening on port " + strconv.Itoa(config.Port))
	if config.CSRF {
		log.Fatal(http.ListenAndServe(":"+strconv.Itoa(config.Port), CSRF(r)))
	} else {
		log.Fatal(http.ListenAndServe(":"+strconv.Itoa(config.Port), r))
	}
}

// RootHandler handles the root route.
// This includes the homepage and the
// file system.
func RootHandler(t *template.Template) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		switch url := r.URL.Path; url {
		case "/":
			if config.CSRF {
				w.Header().Set("X-CSRF-Token", csrf.Token(r))
				err := t.ExecuteTemplate(w, "index.html", map[string]interface{}{
					csrf.TemplateTag: csrf.TemplateField(r),
				})
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
				}
			} else {
				err := t.ExecuteTemplate(w, "index.html", map[string]interface{}{
					csrf.TemplateTag: "",
				})
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
				}
			}
		default:
			if _, err := os.Stat(config.PublicDirectory + url); err != nil {
				w.Header().Add("Content-Type", "text/html; charset=utf-8")
				w.WriteHeader(http.StatusNotFound)
				http.ServeFile(w, r, config.TemplateDirectory+"404.html")
			} else {
				http.ServeFile(w, r, config.PublicDirectory+url)
			}
		}

	}
}

// Handles the ability to view your image.
// It checks if the file exists, and returns
// an appropriate response.
//
// If the image exists, it serves a template
// which is defined in
// TemplateDirectory/view.html.
func ViewHandler(t *template.Template) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		id := vars["id"]
		ext := util.GetFileExtFromDir(id, config.ImageDirectory)
		if ext != "" {
			err := t.ExecuteTemplate(w, "view.html", ViewData{
				Id:       id,
				ImageUrl: config.ImageUrl,
				Ext:      ext,
			})
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				log.Fatal(err)
				return
			}
		} else {
			w.Header().Add("Content-Type", "text/html; charset=utf-8")
			w.WriteHeader(http.StatusNotFound)
			http.ServeFile(w, r, config.TemplateDirectory+"404.html")
		}
	}
}

// Handles the upload route.
// It checks the file, and uploads the
// image to the ImageDirectory defined
// in the config.json
func UploadHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseMultipartForm(30000000)
	file, handler, err := r.FormFile("file")
	if err != nil {
		log.Fatal(err)
		return
	}
	defer file.Close()

	if CheckFileType(handler.Header) == false {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	name := GenerateName(config.ImageNameLength)
	f, err := os.OpenFile(config.ImageDirectory+name+util.GetFileExt(handler.Header["Content-Disposition"][0]), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Fatal(err)
		return
	}

	io.Copy(f, file)
	fStat, err := f.Stat()
	if err != nil {
		log.Fatal(err)
	}
	size := fStat.Size()
	if size > config.MaxFileSize {
		f.Close()
		err = os.RemoveAll(config.ImageDirectory + name + util.GetFileExt(handler.Header["Content-Disposition"][0]))
		if err != nil {
			log.Fatal(err)
		}
		http.Redirect(w, r, "/", http.StatusSeeOther)
	} else {
		http.Redirect(w, r, "/"+name+"/", http.StatusSeeOther)
		f.Close()
	}
}

// CheckFileType checks if the uploaded
// file's extension is allowed.
// It accepts a textproto.MIMEHeader
// which is recieved from the http.Request
// FormFile function, as the handler.Header
func CheckFileType(f textproto.MIMEHeader) bool {
	ext := util.GetFileExt(f["Content-Disposition"][0])

	if !util.Contains(config.AllowedMimeTypes, f["Content-Type"][0]) || !util.Contains(config.AllowedExtensions, ext) {
		return false
	}

	return true
}

// GenerateName generates a name with
// a specified length, and checks if
// the name already exists, using the
// util functions for this.
func GenerateName(n int) string {
	name := util.GenerateName(n)

	for util.CheckExists(name, config.ImageDirectory) {
		name = util.GenerateName(n)
	}

	return name
}
<commit_msg>Update for changed name<commit_after>package main

import (
	"encoding/json"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/textproto"
	"os"
	"strconv"

	"github.com/omar-h/goimage/util"
	"github.com/gorilla/csrf"
	"github.com/gorilla/mux"
)

// Data parsed to the ViewHandler's template
type ViewData struct {
	Id       string
	ImageUrl string
	Ext      string
}

// Parsed JSON config
type Config struct {
	Port              int      `json:"port"`
	Secure            bool     `json:"secure"`
	AuthKey           string   `json:"32-byte-auth-key"`
	AllowedMimeTypes  []string `json:"allowed-mime-types"`
	AllowedExtensions []string `json:"allowed-extensions"`
	ImageNameLength   int      `json:"image-name-length"`
	MaxFileSize       int64    `json:"max-file-size"`
	ImageDirectory    string   `json:"image-directory"`
	TemplateDirectory string   `json:"template-directory"`
	PublicDirectory   string   `json:"public-directory"`
	ImageUrl          string   `json:"image-url"`
	CSRF              bool     `json:"csrf"`
}

var config Config

func main() {
	jsonFiles, err := ioutil.ReadFile("./config.json")
	if err != nil {
		log.Fatal(err)
		return
	}
	json.Unmarshal(jsonFiles, &config)

	CSRF := csrf.Protect(
		[]byte(config.AuthKey),
		csrf.RequestHeader("X-CSRF-Token"),
		csrf.FieldName("_csrf"),
		csrf.Secure(config.Secure),
	)
	r := mux.NewRouter()
	err = os.MkdirAll(config.TemplateDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Template Directory")
	}
	err = os.MkdirAll(config.PublicDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Public Directory")
	}
	err = os.MkdirAll(config.ImageDirectory, 644)
	if err != nil {
		log.Fatal("Unable to create Image Directory")
	}

	templates := template.Must(template.ParseGlob(config.TemplateDirectory + "*.html"))

	r.HandleFunc("/{id}/", ViewHandler(templates)).Methods("GET")
	r.HandleFunc("/upload/", UploadHandler).Methods("POST")
	r.PathPrefix("/").HandlerFunc(RootHandler(templates)).Methods("GET")

	log.Print("Listening on port " + strconv.Itoa(config.Port))
	if config.CSRF {
		log.Fatal(http.ListenAndServe(":"+strconv.Itoa(config.Port), CSRF(r)))
	} else {
		log.Fatal(http.ListenAndServe(":"+strconv.Itoa(config.Port), r))
	}
}

// RootHandler handles the root route.
// This includes the homepage and the
// file system.
func RootHandler(t *template.Template) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		switch url := r.URL.Path; url {
		case "/":
			if config.CSRF {
				w.Header().Set("X-CSRF-Token", csrf.Token(r))
				err := t.ExecuteTemplate(w, "index.html", map[string]interface{}{
					csrf.TemplateTag: csrf.TemplateField(r),
				})
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
				}
			} else {
				err := t.ExecuteTemplate(w, "index.html", map[string]interface{}{
					csrf.TemplateTag: "",
				})
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
				}
			}
		default:
			if _, err := os.Stat(config.PublicDirectory + url); err != nil {
				w.Header().Add("Content-Type", "text/html; charset=utf-8")
				w.WriteHeader(http.StatusNotFound)
				http.ServeFile(w, r, config.TemplateDirectory+"404.html")
			} else {
				http.ServeFile(w, r, config.PublicDirectory+url)
			}
		}

	}
}

// Handles the ability to view your image.
// It checks if the file exists, and returns
// an appropriate response.
//
// If the image exists, it serves a template
// which is defined in
// TemplateDirectory/view.html.
func ViewHandler(t *template.Template) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		id := vars["id"]
		ext := util.GetFileExtFromDir(id, config.ImageDirectory)
		if ext != "" {
			err := t.ExecuteTemplate(w, "view.html", ViewData{
				Id:       id,
				ImageUrl: config.ImageUrl,
				Ext:      ext,
			})
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				log.Fatal(err)
				return
			}
		} else {
			w.Header().Add("Content-Type", "text/html; charset=utf-8")
			w.WriteHeader(http.StatusNotFound)
			http.ServeFile(w, r, config.TemplateDirectory+"404.html")
		}
	}
}

// Handles the upload route.
// It checks the file, and uploads the
// image to the ImageDirectory defined
// in the config.json
func UploadHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseMultipartForm(30000000)
	file, handler, err := r.FormFile("file")
	if err != nil {
		log.Fatal(err)
		return
	}
	defer file.Close()

	if CheckFileType(handler.Header) == false {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	name := GenerateName(config.ImageNameLength)
	f, err := os.OpenFile(config.ImageDirectory+name+util.GetFileExt(handler.Header["Content-Disposition"][0]), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Fatal(err)
		return
	}

	io.Copy(f, file)
	fStat, err := f.Stat()
	if err != nil {
		log.Fatal(err)
	}
	size := fStat.Size()
	if size > config.MaxFileSize {
		f.Close()
		err = os.RemoveAll(config.ImageDirectory + name + util.GetFileExt(handler.Header["Content-Disposition"][0]))
		if err != nil {
			log.Fatal(err)
		}
		http.Redirect(w, r, "/", http.StatusSeeOther)
	} else {
		http.Redirect(w, r, "/"+name+"/", http.StatusSeeOther)
		f.Close()
	}
}

// CheckFileType checks if the uploaded
// file's extension is allowed.
// It accepts a textproto.MIMEHeader
// which is recieved from the http.Request
// FormFile function, as the handler.Header
func CheckFileType(f textproto.MIMEHeader) bool {
	ext := util.GetFileExt(f["Content-Disposition"][0])

	if !util.Contains(config.AllowedMimeTypes, f["Content-Type"][0]) || !util.Contains(config.AllowedExtensions, ext) {
		return false
	}

	return true
}

// GenerateName generates a name with
// a specified length, and checks if
// the name already exists, using the
// util functions for this.
func GenerateName(n int) string {
	name := util.GenerateName(n)

	for util.CheckExists(name, config.ImageDirectory) {
		name = util.GenerateName(n)
	}

	return name
}
<|endoftext|>


<commit_before>// `tyu` is a small command line utility to display basic system informations
// in the terminal.
package main

import (
	ui "github.com/gizak/termui"
)

func main() {
	//init termui
	if err := ui.Init(); err != nil {
		panic(err) //TODO do not panic but manage the error
	}
	defer ui.Close()

	// display physical RAM usage informations
	ramGauge := ui.NewGauge() //TODO try to draw the border around gauge, used, free and total
	ramGauge.BorderLabel = "RAM usage "
	ramGauge.BarColor = ui.ColorBlue
	ramGauge.Width = 39
	ramGauge.Height = 3
	ramGauge.X = 0
	ramGauge.Y = 0

	// display physical swap usage informations
	swapGauge := ui.NewGauge()
	swapGauge.BorderLabel = "Swap usage "
	swapGauge.BarColor = ui.ColorBlue
	swapGauge.Width = 39
	swapGauge.Height = 3
	swapGauge.X = 0
	swapGauge.Y = 3

	cpuGauge := ui.NewGauge()
	cpuGauge.BorderLabel = "CPU usage "
	cpuGauge.BarColor = ui.ColorBlue
	cpuGauge.Width = 39
	cpuGauge.Height = 3
	cpuGauge.X = 0
	cpuGauge.Y = 6

	// display informations about the physical disks
	// TODO only two or 3 physical disk add check and a for loop
	disk := getDiskinfo()
	disk1Gauge := ui.NewGauge()
	disk1Gauge.BorderLabel = disk[0].device + " disk usage "
	disk1Gauge.BarColor = ui.ColorBlue
	disk1Gauge.Width = 39
	disk1Gauge.Height = 3
	disk1Gauge.X = 0
	disk1Gauge.Y = 9
	disk1Gauge.Percent = disk[0].usedPercent

	disk2Gauge := ui.NewGauge()
	disk2Gauge.BorderLabel = disk[1].device + " disk usage " //TODO this fail if only one disk
	disk2Gauge.BarColor = ui.ColorBlue
	disk2Gauge.Width = 39
	disk2Gauge.Height = 3
	disk2Gauge.X = 0
	disk2Gauge.Y = 12
	disk2Gauge.Percent = disk[1].usedPercent

	// display system informations about the host
	host := getHostinfo()
	hostinfo := ui.NewList()
	hostinfo.BorderLabel = "Host "
	hostitems := []string{
		"[Hostname         ](fg-cyan)" + host.hostname,
		"[Domain           ](fg-cyan)" + host.domainname,
		"[OS               ](fg-cyan)" + host.os,
		"[OS version       ](fg-cyan)" + host.osRelease,
		"[Platform         ](fg-cyan)" + host.platform,
		"[Platform version ](fg-cyan)" + host.platformVersion,
		"[Architecture     ](fg-cyan)" + host.arch,
		"[Uptime           ](fg-cyan)",
	}
	hostinfo.Items = hostitems
	hostinfo.Width = 39
	hostinfo.Height = 10
	hostinfo.X = 40
	hostinfo.Y = 0

	// display informations about the CPUs
	cpu := getCPUinfo()
	cpuinfo := ui.NewList()
	cpuinfo.BorderLabel = "CPU "
	cpuinfo.Items = []string{
		"[CPUs        ](fg-cyan)" + cpu.count, //TODO review item names compared to other cpu utilities
		"[Vendor      ](fg-cyan)" + cpu.vendorID,
		"[Model       ](fg-cyan)" + cpu.modelName, //TODO use refreshing rate to display roll long text ?
		"[Frequency   ](fg-cyan)" + cpu.cpuMhz + " Mhz",
		"[Temperature ](fg-cyan)", //TODO
	}
	cpuinfo.Width = 39
	cpuinfo.Height = 7
	cpuinfo.X = 40
	cpuinfo.Y = 10

	// display bios and motherboard informations
	bios := getBIOSinfo()
	biosinfo := ui.NewList()
	biosinfo.BorderLabel = "BIOS "
	biosinfo.Items = []string{
		"[Motherboard ](fg-cyan)" + bios.boardName,
		"[Vendor      ](fg-cyan)" + bios.boardVendor,
		"[BIOS        ](fg-cyan)" + bios.biosVendor,
		"[Version     ](fg-cyan)" + bios.biosVersion + "  " + bios.biosDate,
	}
	biosinfo.Width = 39
	biosinfo.Height = 6
	biosinfo.X = 40
	biosinfo.Y = 17

	// display a quit help text
	quit := ui.NewPar("[Type 'q' to exit](fg-white,bg-blue)")
	quit.Height = 1
	quit.Width = 39
	quit.X = 1
	quit.Y = 23
	quit.Border = false

	// render the dashboard with 26x80 fixed size
	draw := func(t int) {
		// update memory informations
		mem := getMeminfo()
		ramGauge.Percent = mem.ramUsedPercent
		ramGauge.Label = "{{percent}}% - " + mem.ramUsed + "/" + mem.ramTotal + " GiB"
		swapGauge.Percent = mem.swapUsedPercent
		swapGauge.Label = "{{percent}}% - " + mem.swapUsed + "/" + mem.swapTotal + " GiB"

		cpuGauge.Percent = getCPUpercent()

		hostitems[7] = "[Uptime           ](fg-cyan)" + getUptime()

		ui.Render(
			ramGauge,
			swapGauge,
			cpuGauge,
			disk1Gauge, //TODO rename and or stack gauges together
			disk2Gauge, //TODO rename
			hostinfo,
			cpuinfo,
			biosinfo,
			quit,
		)
	}

	// quit on `q` keystroke handler
	ui.Handle("/sys/kbd/q", func(ui.Event) {
		ui.StopLoop()
	})
	// quit on `CTRL+c` keystroke handler
	ui.Handle("/sys/kbd/C-c", func(ui.Event) {
		ui.StopLoop()
	})
	// timer handler to refresh every second
	ui.Handle("/timer/1s", func(e ui.Event) {
		t := e.Data.(ui.EvtTimer)
		draw(int(t.Count))
	})
	ui.Loop()
}

// the dashboard is 26x80
<commit_msg>Update dashboard display<commit_after>// `tyu` is a small command line utility to display basic system informations
// in the terminal.
package main

import (
	ui "github.com/gizak/termui"
)

func main() {
	//init termui
	if err := ui.Init(); err != nil {
		panic(err) //TODO do not panic but manage the error
	}
	defer ui.Close()

	// display physical RAM usage informations
	ramGauge := ui.NewGauge() //TODO try to draw the border around gauge, used, free and total
	ramGauge.BorderLabel = "RAM usage "
	ramGauge.BarColor = ui.ColorBlue
	ramGauge.Width = 39
	ramGauge.Height = 3
	ramGauge.X = 0
	ramGauge.Y = 0

	// display physical swap usage informations
	swapGauge := ui.NewGauge()
	swapGauge.BorderLabel = "Swap usage "
	swapGauge.BarColor = ui.ColorBlue
	swapGauge.Width = 39
	swapGauge.Height = 3
	swapGauge.X = 0
	swapGauge.Y = 3

	cpuGauge := ui.NewGauge()
	cpuGauge.BorderLabel = "CPU usage "
	cpuGauge.BarColor = ui.ColorBlue
	cpuGauge.Width = 39
	cpuGauge.Height = 3
	cpuGauge.X = 0
	cpuGauge.Y = 6

	// display informations about the physical disks
	// TODO only two or 3 physical disk add check and a for loop
	disk := getDiskinfo()
	disk1Gauge := ui.NewGauge()
	disk1Gauge.BorderLabel = disk[0].device + " disk usage "
	disk1Gauge.BarColor = ui.ColorBlue
	disk1Gauge.Width = 39
	disk1Gauge.Height = 3
	disk1Gauge.X = 0
	disk1Gauge.Y = 9
	disk1Gauge.Percent = disk[0].usedPercent

	disk2Gauge := ui.NewGauge()
	disk2Gauge.BorderLabel = disk[1].device + " disk usage " //TODO this fail if only one disk
	disk2Gauge.BarColor = ui.ColorBlue
	disk2Gauge.Width = 39
	disk2Gauge.Height = 3
	disk2Gauge.X = 0
	disk2Gauge.Y = 12
	disk2Gauge.Percent = disk[1].usedPercent

	// display information about the network activity
	netinfo := ui.NewList()
	netinfo.BorderLabel = "Network "
	netitems := []string{
		"[eth0 ](fg-green)", //TODO add a for loop, limit to 2 or 3 interfaces
	}
	netinfo.Items = netitems
	netinfo.Width = 39
	netinfo.Height = 3
	netinfo.X = 0
	netinfo.Y = 15

	// display system informations about the host
	host := getHostinfo()
	hostinfo := ui.NewList()
	hostinfo.BorderLabel = "Host "
	hostitems := []string{
		"[Hostname         ](fg-cyan)" + host.hostname,
		"[Domain           ](fg-cyan)" + host.domainname,
		"[OS               ](fg-cyan)" + host.os,
		"[OS version       ](fg-cyan)" + host.osRelease,
		"[Platform         ](fg-cyan)" + host.platform,
		"[Platform version ](fg-cyan)" + host.platformVersion,
		"[Architecture     ](fg-cyan)" + host.arch,
		"[Uptime           ](fg-cyan)",
	}
	hostinfo.Items = hostitems
	hostinfo.Width = 39
	hostinfo.Height = 10
	hostinfo.X = 40
	hostinfo.Y = 0

	// display informations about the CPUs
	cpu := getCPUinfo()
	cpuinfo := ui.NewList()
	cpuinfo.BorderLabel = "CPU "
	cpuinfo.Items = []string{
		"[CPUs        ](fg-cyan)" + cpu.count, //TODO review item names compared to other cpu utilities
		"[Vendor      ](fg-cyan)" + cpu.vendorID,
		"[Model       ](fg-cyan)" + cpu.modelName, //TODO use refreshing rate to display roll long text ?
		"[Frequency   ](fg-cyan)" + cpu.cpuMhz + " Mhz",
		"[Temperature ](fg-cyan)", //TODO
	}
	cpuinfo.Width = 39
	cpuinfo.Height = 7
	cpuinfo.X = 40
	cpuinfo.Y = 10

	// display bios and motherboard informations
	bios := getBIOSinfo()
	biosinfo := ui.NewList()
	biosinfo.BorderLabel = "BIOS "
	biosinfo.Items = []string{
		"[Motherboard ](fg-cyan)" + bios.boardName,
		"[Vendor      ](fg-cyan)" + bios.boardVendor,
		"[BIOS        ](fg-cyan)" + bios.biosVendor,
		"[Version     ](fg-cyan)" + bios.biosVersion + "  " + bios.biosDate,
	}
	biosinfo.Width = 39
	biosinfo.Height = 6
	biosinfo.X = 40
	biosinfo.Y = 17

	// display a quit help text
	quit := ui.NewPar("[Type 'q' to exit](fg-white,bg-blue)")
	quit.Height = 1
	quit.Width = 39
	quit.X = 1
	quit.Y = 23
	quit.Border = false

	// render the dashboard with 26x80 fixed size
	draw := func(t int) {
		// update memory informations
		mem := getMeminfo()
		ramGauge.Percent = mem.ramUsedPercent
		ramGauge.Label = "{{percent}}% - " + mem.ramUsed + "/" + mem.ramTotal + " GiB"
		swapGauge.Percent = mem.swapUsedPercent
		swapGauge.Label = "{{percent}}% - " + mem.swapUsed + "/" + mem.swapTotal + " GiB"

		net := getNetinfo() //TODO need to calculte before - after see https://github.com/nicolargo/glances/blob/master/glances/plugins/glances_network.py
		netitems[0] = "[eth0 ](fg-green)" + "Up " + net.up + " KiB Down " + net.down + " KiB"

		cpuGauge.Percent = getCPUpercent()

		hostitems[7] = "[Uptime           ](fg-cyan)" + getUptime()

		ui.Render(
			ramGauge,
			swapGauge,
			cpuGauge,
			disk1Gauge, //TODO rename and or stack gauges together
			disk2Gauge, //TODO rename
			netinfo,
			hostinfo,
			cpuinfo,
			biosinfo,
			quit,
		)
	}

	// quit on `q` keystroke handler
	ui.Handle("/sys/kbd/q", func(ui.Event) {
		ui.StopLoop()
	})
	// quit on `CTRL+c` keystroke handler
	ui.Handle("/sys/kbd/C-c", func(ui.Event) {
		ui.StopLoop()
	})
	// timer handler to refresh every second
	ui.Handle("/timer/1s", func(e ui.Event) {
		t := e.Data.(ui.EvtTimer)
		draw(int(t.Count))
	})
	ui.Loop()
}

// the dashboard is 26x80
<|endoftext|>