<commit_before>package main

import (
	. "DNA/client"
	"DNA/common/config"
	"DNA/common/log"
	"DNA/consensus/dbft"
	"DNA/core/ledger"
	"DNA/core/store/ChainStore"
	"DNA/core/transaction"
	"DNA/crypto"
	"DNA/net"
	"DNA/net/httpjsonrpc"
	"fmt"
	"os"
	"runtime"
	"time"
)

const (
	// The number of the CPU cores for parallel optimization,TODO set from config file
	NCPU                   = 4
	DefaultBookKeeperCount = 4
)

var Version string

func init() {
	runtime.GOMAXPROCS(NCPU)
	var path string = "./Log/"
	log.CreatePrintLog(path)
}

func fileExisted(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil || os.IsExist(err)
}

func openLocalClient(name string) Client {
	var c Client

	if fileExisted(name) {
		c = OpenClient(name, []byte("\x12\x34\x56"))
	} else {
		c = CreateClient(name, []byte("\x12\x34\x56"))
	}

	return c
}

func InitBlockChain() ledger.Blockchain {
	blockchain, err := ledger.NewBlockchainWithGenesisBlock()
	if err != nil {
		fmt.Println(err, "  BlockChain generate failed")
	}
	fmt.Println("  BlockChain generate completed. Func test Start...")
	return *blockchain
}

func main() {
	fmt.Printf("Node version: %s\n", Version)
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 0. Client open                                                     ***")
	fmt.Println("//**************************************************************************")
	ledger.DefaultLedger = new(ledger.Ledger)
	ledger.DefaultLedger.Store = ChainStore.NewLedgerStore()
	ledger.DefaultLedger.Store.InitLedgerStore(ledger.DefaultLedger)
	transaction.TxStore = ledger.DefaultLedger.Store
	crypto.SetAlg(crypto.P256R1)
	fmt.Println("  Client set completed. Test Start...")
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 1. Generate [Account]                                              ***")
	fmt.Println("//**************************************************************************")
	var bookKeeperCount uint32 = DefaultBookKeeperCount
	if config.Parameters.BookKeeperCount != 0 {
		bookKeeperCount = config.Parameters.BookKeeperCount
	}
	localclient := OpenClientAndGetAccount(bookKeeperCount)
	if localclient == nil {
		fmt.Println("Can't get local client.")
		os.Exit(1)
	}

	issuer, err := localclient.GetDefaultAccount()
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 2. Set BookKeeper                                                     ***")
	fmt.Println("//**************************************************************************")
	bookKeeper := []*crypto.PubKey{}
	var i uint32
	for i = 0; i < bookKeeperCount; i++ {
		bookKeeper = append(bookKeeper, getBookKeeper(i+1).PublicKey)
	}
	ledger.StandbyBookKeepers = bookKeeper
	fmt.Println("bookKeeper1.PublicKey", issuer.PublicKey)

	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 3. BlockChain init                                                 ***")
	fmt.Println("//**************************************************************************")
	sampleBlockchain := InitBlockChain()
	ledger.DefaultLedger.Blockchain = &sampleBlockchain

	time.Sleep(2 * time.Second)
	neter, noder := net.StartProtocol(issuer.PublicKey)
	httpjsonrpc.RegistRpcNode(noder)
	time.Sleep(20 * time.Second)

	noder.LocalNode().SyncNodeHeight()

	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 5. Start DBFT Services                                             ***")
	fmt.Println("//**************************************************************************")
	dbftServices := dbft.NewDbftService(localclient, "logdbft", neter)
	httpjsonrpc.RegistDbftService(dbftServices)
	go dbftServices.Start()
	time.Sleep(5 * time.Second)
	fmt.Println("DBFT Services start completed.")
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** Init Complete                                                      ***")
	fmt.Println("//**************************************************************************")
	go httpjsonrpc.StartRPCServer()
	go httpjsonrpc.StartLocalServer()

	time.Sleep(2 * time.Second)
	for {
		log.Trace("ledger.DefaultLedger.Blockchain.BlockHeight= ", ledger.DefaultLedger.Blockchain.BlockHeight)
		time.Sleep(dbft.GenBlockTime)
	}
}

func OpenClientAndGetAccount(count uint32) Client {
	clientName := config.Parameters.BookKeeperName
	fmt.Printf("The BookKeeper name is %s\n", clientName)
	if clientName == "" {
		log.Error("BookKeeper name not be set at config.json")
		return nil
	}
	var c []Client
	c = make([]Client, count)
	var i uint32
	for i = 1; i <= count; i++ {
		w := fmt.Sprintf("wallet%d.txt", i)
		if fileExisted(w) {
			c[i-1] = OpenClient(w, []byte("\x12\x34\x56"))
		} else {
			c[i-1] = CreateClient(w, []byte("\x12\x34\x56"))
		}
	}
	var n uint32
	fmt.Sscanf(clientName, "c%d", &n)
	return c[n-1]
}

func getBookKeeper(n uint32) *Account {
	w := fmt.Sprintf("wallet%d.txt", n)
	c := OpenClient(w, []byte("\x12\x34\x56"))
	account, err := c.GetDefaultAccount()
	if err != nil {
		fmt.Println("GetDefaultAccount failed.")
	}
	return account
}
<commit_msg>Set default node count is 4.<commit_after>package main

import (
	. "DNA/client"
	"DNA/common/config"
	"DNA/common/log"
	"DNA/consensus/dbft"
	"DNA/core/ledger"
	"DNA/core/store/ChainStore"
	"DNA/core/transaction"
	"DNA/crypto"
	"DNA/net"
	"DNA/net/httpjsonrpc"
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"
)

const (
	// The number of the CPU cores for parallel optimization,TODO set from config file
	NCPU                   = 4
	DefaultBookKeeperCount = 4
)

var Version string

func init() {
	runtime.GOMAXPROCS(NCPU)
	var path string = "./Log/"
	log.CreatePrintLog(path)
}

func fileExisted(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil || os.IsExist(err)
}

func openLocalClient(name string) Client {
	var c Client

	if fileExisted(name) {
		c = OpenClient(name, []byte("\x12\x34\x56"))
	} else {
		c = CreateClient(name, []byte("\x12\x34\x56"))
	}

	return c
}

func InitBlockChain() ledger.Blockchain {
	blockchain, err := ledger.NewBlockchainWithGenesisBlock()
	if err != nil {
		fmt.Println(err, "  BlockChain generate failed")
	}
	fmt.Println("  BlockChain generate completed. Func test Start...")
	return *blockchain
}

func main() {
	fmt.Printf("Node version: %s\n", Version)
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 0. Client open                                                     ***")
	fmt.Println("//**************************************************************************")
	ledger.DefaultLedger = new(ledger.Ledger)
	ledger.DefaultLedger.Store = ChainStore.NewLedgerStore()
	ledger.DefaultLedger.Store.InitLedgerStore(ledger.DefaultLedger)
	transaction.TxStore = ledger.DefaultLedger.Store
	crypto.SetAlg(crypto.P256R1)
	fmt.Println("  Client set completed. Test Start...")
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 1. Generate [Account]                                              ***")
	fmt.Println("//**************************************************************************")

	localclient := OpenClientAndGetAccount()
	if localclient == nil {
		fmt.Println("Can't get local client.")
		os.Exit(1)
	}

	issuer, err := localclient.GetDefaultAccount()
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 2. Set BookKeeper                                                     ***")
	fmt.Println("//**************************************************************************")
	bookKeeper := []*crypto.PubKey{}
	var i uint32
	for i = 0; i < DefaultBookKeeperCount; i++ {
		bookKeeper = append(bookKeeper, getBookKeeper(i+1).PublicKey)
	}
	ledger.StandbyBookKeepers = bookKeeper
	fmt.Println("bookKeeper1.PublicKey", issuer.PublicKey)

	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 3. BlockChain init                                                 ***")
	fmt.Println("//**************************************************************************")
	sampleBlockchain := InitBlockChain()
	ledger.DefaultLedger.Blockchain = &sampleBlockchain

	time.Sleep(2 * time.Second)
	neter, noder := net.StartProtocol(issuer.PublicKey)
	httpjsonrpc.RegistRpcNode(noder)
	time.Sleep(20 * time.Second)

	noder.LocalNode().SyncNodeHeight()

	fmt.Println("//**************************************************************************")
	fmt.Println("//*** 5. Start DBFT Services                                             ***")
	fmt.Println("//**************************************************************************")
	dbftServices := dbft.NewDbftService(localclient, "logdbft", neter)
	httpjsonrpc.RegistDbftService(dbftServices)
	go dbftServices.Start()
	time.Sleep(5 * time.Second)
	fmt.Println("DBFT Services start completed.")
	fmt.Println("//**************************************************************************")
	fmt.Println("//*** Init Complete                                                      ***")
	fmt.Println("//**************************************************************************")
	go httpjsonrpc.StartRPCServer()
	go httpjsonrpc.StartLocalServer()

	time.Sleep(2 * time.Second)
	for {
		log.Trace("ledger.DefaultLedger.Blockchain.BlockHeight= ", ledger.DefaultLedger.Blockchain.BlockHeight)
		time.Sleep(dbft.GenBlockTime)
	}
}

func clientIsDefaultBookKeeper(clientName string) bool {
	var i uint32
	for i = 0; i < DefaultBookKeeperCount; i++ {
		c := fmt.Sprintf("c%d", i+1)
		if strings.Compare(clientName, c) == 0 {
			return true
		}
	}
	return false
}

func OpenClientAndGetAccount() Client {
	clientName := config.Parameters.BookKeeperName
	fmt.Printf("The BookKeeper name is %s\n", clientName)
	if clientName == "" {
		log.Error("BookKeeper name not be set at config.json")
		return nil
	}
	isDefaultBookKeeper := clientIsDefaultBookKeeper(clientName)
	var c []Client
	if isDefaultBookKeeper == true {
		c = make([]Client, DefaultBookKeeperCount)
	} else {
		c = make([]Client, DefaultBookKeeperCount+1)
	}
	var i uint32
	for i = 1; i <= DefaultBookKeeperCount; i++ {
		w := fmt.Sprintf("wallet%d.txt", i)
		if fileExisted(w) {
			c[i-1] = OpenClient(w, []byte("\x12\x34\x56"))
		} else {
			c[i-1] = CreateClient(w, []byte("\x12\x34\x56"))
		}
	}
	var n uint32
	fmt.Sscanf(clientName, "c%d", &n)
	if isDefaultBookKeeper == true {
		return c[n-1]
	}
	if isDefaultBookKeeper == false {
		w := fmt.Sprintf("wallet%d.txt", n)
		if fileExisted(w) {
			c[DefaultBookKeeperCount] = OpenClient(w, []byte("\x12\x34\x56"))
		} else {
			c[DefaultBookKeeperCount] = CreateClient(w, []byte("\x12\x34\x56"))
		}
	}
	return c[DefaultBookKeeperCount]
}

func getBookKeeper(n uint32) *Account {
	w := fmt.Sprintf("wallet%d.txt", n)
	c := OpenClient(w, []byte("\x12\x34\x56"))
	account, err := c.GetDefaultAccount()
	if err != nil {
		fmt.Println("GetDefaultAccount failed.")
	}
	return account
}
<|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.4.1"

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", "", "schema name for drivers that support it (default psql: public, mssql: dbo)")
	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.SetDefault("mssql.sslmode", "true")
	viper.SetDefault("mssql.port", "1433")

	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)
		}

		if len(cmdConfig.Schema) == 0 {
			cmdConfig.Schema = "public"
		}

		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())
		}
	}

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

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

		if cmdConfig.MSSQL.Port == 0 {
			cmdConfig.MSSQL.Port = 1433
			viper.Set("mssql.port", cmdConfig.MSSQL.Port)
		}

		if len(cmdConfig.Schema) == 0 {
			cmdConfig.Schema = "dbo"
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.MSSQL.User, "mssql.user"),
			vala.StringNotEmpty(cmdConfig.MSSQL.Host, "mssql.host"),
			vala.Not(vala.Equals(cmdConfig.MSSQL.Port, 0, "mssql.port")),
			vala.StringNotEmpty(cmdConfig.MSSQL.DBName, "mssql.dbname"),
			vala.StringNotEmpty(cmdConfig.MSSQL.SSLMode, "mssql.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>Bump version 2.5.0<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.5.0"

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", "", "schema name for drivers that support it (default psql: public, mssql: dbo)")
	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.SetDefault("mssql.sslmode", "true")
	viper.SetDefault("mssql.port", "1433")

	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)
		}

		if len(cmdConfig.Schema) == 0 {
			cmdConfig.Schema = "public"
		}

		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())
		}
	}

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

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

		if cmdConfig.MSSQL.Port == 0 {
			cmdConfig.MSSQL.Port = 1433
			viper.Set("mssql.port", cmdConfig.MSSQL.Port)
		}

		if len(cmdConfig.Schema) == 0 {
			cmdConfig.Schema = "dbo"
		}

		err = vala.BeginValidation().Validate(
			vala.StringNotEmpty(cmdConfig.MSSQL.User, "mssql.user"),
			vala.StringNotEmpty(cmdConfig.MSSQL.Host, "mssql.host"),
			vala.Not(vala.Equals(cmdConfig.MSSQL.Port, 0, "mssql.port")),
			vala.StringNotEmpty(cmdConfig.MSSQL.DBName, "mssql.dbname"),
			vala.StringNotEmpty(cmdConfig.MSSQL.SSLMode, "mssql.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 (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	arg "github.com/alexflint/go-arg"
	"github.com/joho/godotenv"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/template"
	"github.com/prometheus/common/model"
	"github.com/tucnak/telebot"
)

const (
	commandStart = "/start"
	commandStop  = "/stop"
	commandHelp  = "/help"
	commandUsers = "/users"

	commandStatus     = "/status"
	commandAlerts     = "/alerts"
	commandSilences   = "/silences"
	commandSilenceAdd = "/silence_add"
	commandSilence    = "/silence"
	commandSilenceDel = "/silence_del"

	responseStart = "Hey, %s! I will now keep you up to date!\n" + commandHelp
	responseStop  = "Alright, %s! I won't talk to you again.\n" + commandHelp
	responseHelp  = `
I'm a Prometheus AlertManager bot for Telegram. I will notify you about alerts.
You can also ask me about my ` + commandStatus + `, ` + commandAlerts + ` & ` + commandSilences + `

Available commands:
` + commandStart + ` - Subscribe for alerts.
` + commandStop + ` - Unsubscribe for alerts.
` + commandStatus + ` - Print the current status.
` + commandAlerts + ` - List all alerts.
` + commandSilences + ` - List all silences.
`
)

var users map[int]telebot.User

type Config struct {
	AlertmanagerURL string `arg:"env:ALERTMANAGER_URL"`
	TelegramToken   string `arg:"env:TELEGRAM_TOKEN"`
	TelegramAdmin   int    `arg:"env:TELEGRAM_ADMIN"`
}

func main() {
	log.Println("starting...")

	if err := godotenv.Load(); err != nil {
		log.Println(err)
	}

	// initialize Config{} with default values
	c := Config{AlertmanagerURL: "http://localhost:9093"}
	arg.MustParse(&c)

	bot, err := telebot.NewBot(c.TelegramToken)
	if err != nil {
		log.Fatalln(err)
	}

	users = make(map[int]telebot.User)
	messages := make(chan telebot.Message, 100)
	bot.Listen(messages, 1*time.Second)

	go HTTPListenAndServe(bot)

	for message := range messages {
		if message.Sender.ID != c.TelegramAdmin {
			log.Printf("dropped message from unallowed sender: %s(%d)", message.Sender.Username, message.Sender.ID)
			continue
		}

		switch message.Text {
		case commandStart:
			bot.SendMessage(message.Chat, fmt.Sprintf(responseStart, message.Sender.FirstName), nil)
			users[message.Sender.ID] = message.Sender
			log.Printf("User %s(%d) subscribed", message.Sender.Username, message.Sender.ID)
		case commandStop:
			bot.SendMessage(message.Chat, fmt.Sprintf(responseStop, message.Sender.FirstName), nil)
			delete(users, message.Sender.ID)
			log.Printf("User %s(%d) unsubscribed", message.Sender.Username, message.Sender.ID)
		case commandHelp:
			bot.SendMessage(message.Chat, responseHelp, nil)
		case commandUsers:
			bot.SendMessage(message.Chat, fmt.Sprintf("Currently %d users are subscribed.", len(users)), nil)
		case commandAlerts:
			alerts, err := listAlerts(c)
			if err != nil {
				bot.SendMessage(message.Chat, fmt.Sprintf("failed to list alerts... %v", err), nil)
			}

			var out string
			for _, a := range alerts {
				out = out + Message(a) + "\n"
			}

			bot.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
		}
	}
}

func HTTPListenAndServe(bot *telebot.Bot) {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var webhook notify.WebhookMessage

		var buf bytes.Buffer
		tee := io.TeeReader(r.Body, &buf)
		defer r.Body.Close()

		decoder := json.NewDecoder(tee)
		if err := decoder.Decode(&webhook); err != nil {
			log.Printf("failed to decode webhook message: %v\n", err)
		}

		body, err := ioutil.ReadAll(&buf)
		if err != nil {
			log.Printf("failed to read from request.Body for logging: %v", err)
		}
		log.Println(string(body))

		for _, alert := range webhook.Alerts {
			var out string
			out = out + Message(alert) + "\n"

			for _, user := range users {
				bot.SendMessage(user, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
			}

		}

		w.WriteHeader(http.StatusOK)
	})

	log.Fatalln(http.ListenAndServe(":8080", nil))
}

func Message(a template.Alert) string {
	if a.Status == "" {
		if a.EndsAt.IsZero() {
			a.Status = string(model.AlertFiring)
		} else {
			a.Status = string(model.AlertResolved)
		}
	}

	status := a.Status
	switch status {
	case string(model.AlertFiring):
		status = "🔥 *" + strings.ToUpper(status) + "* 🔥"
	case string(model.AlertResolved):
		status = "*" + strings.ToUpper(status) + "*"
	}

	return fmt.Sprintf(
		"%s\n*%s* (%s)\n%s\n",
		status,
		a.Labels["alertname"],
		a.Annotations["summary"],
		a.Annotations["description"],
	)
}

type alertResponse struct {
	Status string           `json:"status"`
	Alerts []template.Alert `json:"data,omitempty"`
}

func listAlerts(c Config) ([]template.Alert, error) {
	resp, err := http.Get(c.AlertmanagerURL + "/api/v1/alerts")
	if err != nil {
		return nil, err
	}

	var alertResponse alertResponse
	dec := json.NewDecoder(resp.Body)
	dec.Decode(&alertResponse)

	return alertResponse.Alerts, err
}
<commit_msg>Quickly add a first draft of status command<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	arg "github.com/alexflint/go-arg"
	"github.com/joho/godotenv"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/template"
	"github.com/prometheus/common/model"
	"github.com/tucnak/telebot"
)

const (
	commandStart = "/start"
	commandStop  = "/stop"
	commandHelp  = "/help"
	commandUsers = "/users"

	commandStatus     = "/status"
	commandAlerts     = "/alerts"
	commandSilences   = "/silences"
	commandSilenceAdd = "/silence_add"
	commandSilence    = "/silence"
	commandSilenceDel = "/silence_del"

	responseStart = "Hey, %s! I will now keep you up to date!\n" + commandHelp
	responseStop  = "Alright, %s! I won't talk to you again.\n" + commandHelp
	responseHelp  = `
I'm a Prometheus AlertManager bot for Telegram. I will notify you about alerts.
You can also ask me about my ` + commandStatus + `, ` + commandAlerts + ` & ` + commandSilences + `

Available commands:
` + commandStart + ` - Subscribe for alerts.
` + commandStop + ` - Unsubscribe for alerts.
` + commandStatus + ` - Print the current status.
` + commandAlerts + ` - List all alerts.
` + commandSilences + ` - List all silences.
`
)

var users map[int]telebot.User

type Config struct {
	AlertmanagerURL string `arg:"env:ALERTMANAGER_URL"`
	TelegramToken   string `arg:"env:TELEGRAM_TOKEN"`
	TelegramAdmin   int    `arg:"env:TELEGRAM_ADMIN"`
}

func main() {
	log.Println("starting...")

	if err := godotenv.Load(); err != nil {
		log.Println(err)
	}

	// initialize Config{} with default values
	c := Config{AlertmanagerURL: "http://localhost:9093"}
	arg.MustParse(&c)

	bot, err := telebot.NewBot(c.TelegramToken)
	if err != nil {
		log.Fatalln(err)
	}

	users = make(map[int]telebot.User)
	messages := make(chan telebot.Message, 100)
	bot.Listen(messages, 1*time.Second)

	go HTTPListenAndServe(bot)

	for message := range messages {
		if message.Sender.ID != c.TelegramAdmin {
			log.Printf("dropped message from unallowed sender: %s(%d)", message.Sender.Username, message.Sender.ID)
			continue
		}

		switch message.Text {
		case commandStart:
			bot.SendMessage(message.Chat, fmt.Sprintf(responseStart, message.Sender.FirstName), nil)
			users[message.Sender.ID] = message.Sender
			log.Printf("User %s(%d) subscribed", message.Sender.Username, message.Sender.ID)
		case commandStop:
			bot.SendMessage(message.Chat, fmt.Sprintf(responseStop, message.Sender.FirstName), nil)
			delete(users, message.Sender.ID)
			log.Printf("User %s(%d) unsubscribed", message.Sender.Username, message.Sender.ID)
		case commandHelp:
			bot.SendMessage(message.Chat, responseHelp, nil)
		case commandUsers:
			bot.SendMessage(message.Chat, fmt.Sprintf("Currently %d users are subscribed.", len(users)), nil)
		case commandStatus:
			s, err := status(c)
			if err != nil {
				bot.SendMessage(message.Chat, fmt.Sprintf("failed to get status... %v", err), nil)
			}
			bot.SendMessage(message.Chat, fmt.Sprintf("Version: %s\nUptime: %s", s.Data.VersionInfo.Version, s.Data.Uptime), nil)
		case commandAlerts:
			alerts, err := listAlerts(c)
			if err != nil {
				bot.SendMessage(message.Chat, fmt.Sprintf("failed to list alerts... %v", err), nil)
			}

			var out string
			for _, a := range alerts {
				out = out + Message(a) + "\n"
			}

			bot.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
		}
	}
}

func HTTPListenAndServe(bot *telebot.Bot) {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var webhook notify.WebhookMessage

		var buf bytes.Buffer
		tee := io.TeeReader(r.Body, &buf)
		defer r.Body.Close()

		decoder := json.NewDecoder(tee)
		if err := decoder.Decode(&webhook); err != nil {
			log.Printf("failed to decode webhook message: %v\n", err)
		}

		body, err := ioutil.ReadAll(&buf)
		if err != nil {
			log.Printf("failed to read from request.Body for logging: %v", err)
		}
		log.Println(string(body))

		for _, alert := range webhook.Alerts {
			var out string
			out = out + Message(alert) + "\n"

			for _, user := range users {
				bot.SendMessage(user, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
			}

		}

		w.WriteHeader(http.StatusOK)
	})

	log.Fatalln(http.ListenAndServe(":8080", nil))
}

func Message(a template.Alert) string {
	if a.Status == "" {
		if a.EndsAt.IsZero() {
			a.Status = string(model.AlertFiring)
		} else {
			a.Status = string(model.AlertResolved)
		}
	}

	status := a.Status
	switch status {
	case string(model.AlertFiring):
		status = "🔥 *" + strings.ToUpper(status) + "* 🔥"
	case string(model.AlertResolved):
		status = "*" + strings.ToUpper(status) + "*"
	}

	return fmt.Sprintf(
		"%s\n*%s* (%s)\n%s\n",
		status,
		a.Labels["alertname"],
		a.Annotations["summary"],
		a.Annotations["description"],
	)
}

type alertResponse struct {
	Status string           `json:"status"`
	Alerts []template.Alert `json:"data,omitempty"`
}

func listAlerts(c Config) ([]template.Alert, error) {
	resp, err := http.Get(c.AlertmanagerURL + "/api/v1/alerts")
	if err != nil {
		return nil, err
	}

	var alertResponse alertResponse
	dec := json.NewDecoder(resp.Body)
	defer resp.Body.Close()
	dec.Decode(&alertResponse)

	return alertResponse.Alerts, err
}

type statusResponse struct {
	Status string `json:"status"`
	Data   struct {
		Uptime      time.Time `json:"uptime"`
		VersionInfo struct {
			Branch    string `json:"branch"`
			BuildDate string `json:"buildDate"`
			BuildUser string `json:"buildUser"`
			GoVersion string `json:"goVersion"`
			Revision  string `json:"revision"`
			Version   string `json:"version"`
		} `json:"versionInfo"`
	} `json:"data"`
}

func status(c Config) (statusResponse, error) {
	var statusResponse statusResponse

	resp, err := http.Get(c.AlertmanagerURL + "/api/v1/status")
	if err != nil {
		return statusResponse, err
	}

	dec := json.NewDecoder(resp.Body)
	defer resp.Body.Close()
	dec.Decode(&statusResponse)

	return statusResponse, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	"github.com/prodhe/slides/parse"
)

func main() {
	flag.Parse()

	var f *os.File
	var err error
	var path string

	if flag.NArg() < 1 || flag.Arg(0) == "-" {
		f = os.Stdin
	} else {
		path = flag.Arg(0)
		f, err = os.Open(path)
		if err != nil {
			fmt.Printf("%s: %v\n", path, err)
			os.Exit(1)
		}
	}
	reader := bufio.NewReader(f)
	input, err := ioutil.ReadAll(reader)
	if err != nil {
		fmt.Printf("ReadAll error: %v\n", err)
		os.Exit(1)
	}
	f.Close()

	p := parse.NewParser(path, string(input))

	data, err := p.Parse()
	if err != nil {
		fmt.Printf("parse error: %s", err)
		os.Exit(1)
	}

	http.Handle("/", slides(toHtml(data)))
	fmt.Println("Slides are available at http://localhost:3001/")
	err = http.ListenAndServe(":3001", nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func slides(data string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, data)
	})
}

func toHtml(data string) string {
	result := `<!doctype html>
<html>
<head>
<title>Slides</title>
{{style}}
</head>
<body>
{{data}}
</body>
</html>
`

	style := `<style type="text/css">
* { border: 0; margin: 0; padding: 0; box-sizing: border-box; }
body { background-color: #ffffea; }
p {
	width: 100vw;
	height: 100vh;
	display: flex;
	align-items: center;
	justify-content: center;
	font-size: 24pt;
	font-family: monospace;
	line-height: 1.7;
	text-align: left;
	padding: 1em;
}
</style>`

	result = strings.Replace(result, "{{style}}", style, -1)
	result = strings.Replace(result, "{{data}}", data, -1)

	return result
}
<commit_msg>web: Make slides scroll sideways.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	"github.com/prodhe/slides/parse"
)

func main() {
	flag.Parse()

	var f *os.File
	var err error
	var path string

	if flag.NArg() < 1 || flag.Arg(0) == "-" {
		f = os.Stdin
	} else {
		path = flag.Arg(0)
		f, err = os.Open(path)
		if err != nil {
			fmt.Printf("%s: %v\n", path, err)
			os.Exit(1)
		}
	}
	reader := bufio.NewReader(f)
	input, err := ioutil.ReadAll(reader)
	if err != nil {
		fmt.Printf("ReadAll error: %v\n", err)
		os.Exit(1)
	}
	f.Close()

	p := parse.NewParser(path, string(input))

	data, err := p.Parse()
	if err != nil {
		fmt.Printf("parse error: %s", err)
		os.Exit(1)
	}

	http.Handle("/", slides(toHtml(data)))
	fmt.Println("Slides are available at http://localhost:3001/")
	err = http.ListenAndServe(":3001", nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func slides(data string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, data)
	})
}

func toHtml(data string) string {
	result := `<!doctype html>
<html>
<head>
<title>Slides</title>
{{style}}
</head>
<body>
<div id="slides">
{{data}}
</div>
</body>
</html>
`

	style := `<style type="text/css">
* { border: 0; margin: 0; padding: 0; box-sizing: border-box; }
body { background-color: #ffffea; }
#slides {
	display: flex;
	flex-flow: row nowrap;
}
p {
	flex: none;
	width: 100vw;
	height: 100vh;
	display: flex;
	align-items: center;
	justify-content: center;
	font-size: 24pt;
	font-family: monospace;
	line-height: 1.7;
	text-align: left;
	padding: 1em;
}
</style>`

	result = strings.Replace(result, "{{style}}", style, -1)
	result = strings.Replace(result, "{{data}}", data, -1)

	return result
}
<|endoftext|>


<commit_before>// main
package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/sloonz/go-iconv"
	"golang.org/x/crypto/ssh/terminal"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"syscall"
)

const (
	EVENT_ID_START            = "go_event('"
	EVENT_ID_END              = "', "
	CONTENT_START             = "style=\"padding:7px 5px 7px 5px; line-height:1.6;\" bgcolor=\"#FFFFFF\" align=\"left\">"
	CONTENT_END               = "	"
	TYPE_DAILY                = 1
	TYPE_CONTENT              = 2
	TYPE_TIP                  = 3
	TYPE_GOODS                = 4
	TYPE_COMPANY              = 5
	CURRENT_EVENT_ID_POSITION = "F9F5C6"
	TITLE_START               = ""
	TITLE_END                 = ""
	DAILY                     = "icon_daily.gif"
	TIP_START                 = "style=\"padding:5px 5px 5px 5px; line-height:1.6;\" bgcolor=\"#FFFFFF\" align=\"left\">"
	TIP_END                   = "	"
	GOODS_START               = "3px 5px 0px 5px; word-break: break-all\">"
	GOODS_END                 = "	"
)

type FilterEntry struct {
	typeValue int
	regex     string
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("필터 파일명을 입력해주세요.")
		return
	}
	// 필터 파일 읽어와서 배열로 만듦
	filterFile, error := os.Open(os.Args[1])
	if error != nil {
		fmt.Println("필터 파일이 없습니다.")
		return
	}
	defer filterFile.Close()
	scanner := bufio.NewScanner(filterFile)
	typeOrder := true
	var typeValue int
	var filter []FilterEntry
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if len(line) > 0 && line[0] != '#' {
			if typeOrder {
				typeValue, error = strconv.Atoi(line)
				if error != nil {
					fmt.Println("[필터 파일 오류]타입으로 변환하지 못했습니다. 해당 줄 :", line)
					return
				}
				if typeValue != TYPE_DAILY {
					typeOrder = false
				} else {
					filter = append(filter, FilterEntry{typeValue, ""})
				}
			} else {
				filter = append(filter, FilterEntry{typeValue, line})
				typeOrder = true
			}
		}
	}
	if !typeOrder {
		fmt.Println("[필터 파일 오류]정규표현식이 입력되지 않았습니다.")
		return
	}

	// ID, Password 입력 받음(id, password)
	var id string

	fmt.Print("ID : ")
	fmt.Scan(&id)
	fmt.Print("Password : ")
	passwordByte, _ := terminal.ReadPassword(syscall.Stdin)
	password := string(passwordByte)
	fmt.Println("")

	// 로그인 시도
	response, _ := http.PostForm("http://www.e-pass.co.kr/member/login_proc.asp",
		url.Values{"tm_userid": {url.QueryEscape(id)}, "tm_pwd": {url.QueryEscape(password)}})
	defer response.Body.Close()
	body, _ := ioutil.ReadAll(response.Body)

	bodyString := string(body)
	if strings.Index(bodyString, "cation.href") < 0 {
		fmt.Println("로그인에 실패했습니다.")
		return
	}

	// 쿠키 저장
	cookies := response.Cookies()

	// 이벤트 전체 페이지에서 첫 번째 이벤트 ID 가져오기
	client := http.Client{}
	request := cookieAddedRequest("http://e-pass.co.kr/event/all.asp", cookies)
	response, _ = client.Do(request)
	body, _ = ioutil.ReadAll(response.Body)
	html, _ := iconv.Conv(string(body), "utf-8", "euc-kr")
	index := strings.Index(html, EVENT_ID_START)
	if index < 0 {
		fmt.Println("이벤트가 없습니다.")
		return
	}
	eventId := html[index+len(EVENT_ID_START):]
	eventId = eventId[:strings.Index(eventId, EVENT_ID_END)]
	MF := 0
	for {
		url := "http://e-pass.co.kr/event/all_info.asp?" +
			"InNo=" + eventId + "&" +
			"MF=" + strconv.Itoa(MF)
		request = cookieAddedRequest(url, cookies)
		response, _ = client.Do(request)
		body, _ = ioutil.ReadAll(response.Body)
		html, error = iconv.Conv(string(body), "utf-8", "euc-kr")
		// 가끔 인코딩이 안될 때가 있음
		if error != nil {
			fmt.Println("인코딩에 실패했습니다. URL :", url)
			MF++
			eventId = nextEventId(string(body))
			if len(eventId) == 0 {
				break
			}
			continue
		}
		title := html[strings.Index(html, TITLE_START)+len(TITLE_START):]
		title = title[:strings.Index(title, TITLE_END)]
		content := ""
		tip := ""
		tipFind := false
		filtered := false
		company := ""
		var goods []string
		for _, entry := range filter {
			switch entry.typeValue {
			// 매일 응모
			case TYPE_DAILY:
				if strings.Index(html, DAILY) >= 0 {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 행사 내용
			case TYPE_CONTENT:
				if len(content) == 0 {
					content = html[strings.Index(html, CONTENT_START)+len(CONTENT_START):]
					content = strings.TrimSpace(content[:strings.Index(content, CONTENT_END)])
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				if regex.MatchString(content) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 응모 요령
			case TYPE_TIP:
				if !tipFind {
					index = strings.Index(html, TIP_START)
					if index >= 0 {
						tip = html[strings.Index(html, TIP_START)+len(TIP_START):]
						tip = strings.TrimSpace(tip[:strings.Index(tip, TIP_END)])
					}
					tipFind = true
				}
				if len(tip) > 0 {
					regex, error := regexp.Compile(entry.regex)
					if error != nil {
						fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
						return
					}
					if regex.MatchString(tip) {
						deleteEvent(eventId, cookies)
						filtered = true
					}
				}
			// 경품
			case TYPE_GOODS:
				if len(goods) == 0 {
					for goodsString := html; ; {
						index = strings.Index(goodsString, GOODS_START)
						if index < 0 {
							break
						}
						goodsString = goodsString[index+len(GOODS_START):]
						oneGoods := strings.TrimSpace(goodsString[:strings.Index(goodsString, GOODS_END)])
						goods = append(goods, oneGoods)
					}
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				matchCount := 0
				for _, oneGoods := range goods {
					if regex.MatchString(oneGoods) {
						matchCount++
					}
				}
				if matchCount == len(goods) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 주최사
			case TYPE_COMPANY:
				if len(company) == 0 {
					company = title[strings.Index(title, "[")+1 : strings.Index(title, "]")]
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				if regex.MatchString(company) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			}
			if filtered {
				break
			}
		}
		fmt.Print(title + " => ")
		if filtered {
			fmt.Println("삭제")
		} else {
			fmt.Println("통과")
			MF++
		}
		eventId = nextEventId(html)
		if len(eventId) == 0 {
			break
		}
	}
}

func cookieAddedRequest(url string, cookies []*http.Cookie) *http.Request {
	request, _ := http.NewRequest("GET", url, nil)

	for _, cookie := range cookies {
		request.AddCookie(cookie)
	}

	return request
}

func deleteEvent(eventId string, cookies []*http.Cookie) {
	client := http.Client{}
	params := url.Values{"ApplyType": {"1,0"}, "InNo": {eventId}}
	request, _ := http.NewRequest("POST", "http://e-pass.co.kr/user/event_apply_detail.asp",
		bytes.NewBufferString(params.Encode()))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	for _, cookie := range cookies {
		request.AddCookie(cookie)
	}
	client.Do(request)
}

func nextEventId(html string) string {
	index := strings.Index(html, CURRENT_EVENT_ID_POSITION)
	if index < 0 {
		return ""
	}
	eventId := html[index:]
	eventId = eventId[strings.Index(eventId, EVENT_ID_START)+len(EVENT_ID_START):]
	eventId = eventId[strings.Index(eventId, EVENT_ID_START)+len(EVENT_ID_START):]
	index = strings.Index(eventId, EVENT_ID_START)
	if index < 0 {
		return ""
	}
	eventId = eventId[index+len(EVENT_ID_START):]
	eventId = eventId[:strings.Index(eventId, EVENT_ID_END)]

	return eventId
}
<commit_msg>발표방법 추가<commit_after>// main
package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/sloonz/go-iconv"
	"golang.org/x/crypto/ssh/terminal"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"syscall"
)

const (
	EVENT_ID_START            = "go_event('"
	EVENT_ID_END              = "', "
	CONTENT_START             = "style=\"padding:7px 5px 7px 5px; line-height:1.6;\" bgcolor=\"#FFFFFF\" align=\"left\">"
	CONTENT_END               = "	"
	TYPE_DAILY                = 1
	TYPE_CONTENT              = 2
	TYPE_TIP                  = 3
	TYPE_GOODS                = 4
	TYPE_COMPANY              = 5
	TYPE_ANNOUNCE_METHOD      = 6
	CURRENT_EVENT_ID_POSITION = "F9F5C6"
	TITLE_START               = ""
	TITLE_END                 = ""
	DAILY                     = "icon_daily.gif"
	TIP_START                 = "style=\"padding:5px 5px 5px 5px; line-height:1.6;\" bgcolor=\"#FFFFFF\" align=\"left\">"
	TIP_END                   = "	"
	GOODS_START               = "3px 5px 0px 5px; word-break: break-all\">"
	GOODS_END                 = "	"
	ANNOUNCE_METHOD_START     = "style=\"padding:3px 0px 0px 5px;line-height:1.7;\" align=\"left\" bgcolor=\"#FFFFFF\">"
	ANNOUNCE_METHOD_END       = "	"
)

type FilterEntry struct {
	typeValue int
	regex     string
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("필터 파일명을 입력해주세요.")
		return
	}
	// 필터 파일 읽어와서 배열로 만듦
	filterFile, error := os.Open(os.Args[1])
	if error != nil {
		fmt.Println("필터 파일이 없습니다.")
		return
	}
	defer filterFile.Close()
	scanner := bufio.NewScanner(filterFile)
	typeOrder := true
	var typeValue int
	var filter []FilterEntry
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if len(line) > 0 && line[0] != '#' {
			if typeOrder {
				typeValue, error = strconv.Atoi(line)
				if error != nil {
					fmt.Println("[필터 파일 오류]타입으로 변환하지 못했습니다. 해당 줄 :", line)
					return
				}
				if typeValue != TYPE_DAILY {
					typeOrder = false
				} else {
					filter = append(filter, FilterEntry{typeValue, ""})
				}
			} else {
				filter = append(filter, FilterEntry{typeValue, line})
				typeOrder = true
			}
		}
	}
	if !typeOrder {
		fmt.Println("[필터 파일 오류]정규표현식이 입력되지 않았습니다.")
		return
	}

	// ID, Password 입력 받음(id, password)
	var id string

	fmt.Print("ID : ")
	fmt.Scan(&id)
	fmt.Print("Password : ")
	passwordByte, _ := terminal.ReadPassword(syscall.Stdin)
	password := string(passwordByte)
	fmt.Println("")

	// 로그인 시도
	response, _ := http.PostForm("http://www.e-pass.co.kr/member/login_proc.asp",
		url.Values{"tm_userid": {url.QueryEscape(id)}, "tm_pwd": {url.QueryEscape(password)}})
	defer response.Body.Close()
	body, _ := ioutil.ReadAll(response.Body)

	bodyString := string(body)
	if strings.Index(bodyString, "cation.href") < 0 {
		fmt.Println("로그인에 실패했습니다.")
		return
	}

	// 쿠키 저장
	cookies := response.Cookies()

	// 이벤트 전체 페이지에서 첫 번째 이벤트 ID 가져오기
	client := http.Client{}
	request := cookieAddedRequest("http://e-pass.co.kr/event/all.asp", cookies)
	response, _ = client.Do(request)
	body, _ = ioutil.ReadAll(response.Body)
	html, _ := iconv.Conv(string(body), "utf-8", "euc-kr")
	index := strings.Index(html, EVENT_ID_START)
	if index < 0 {
		fmt.Println("이벤트가 없습니다.")
		return
	}
	eventId := html[index+len(EVENT_ID_START):]
	eventId = eventId[:strings.Index(eventId, EVENT_ID_END)]
	MF := 0
	for {
		url := "http://e-pass.co.kr/event/all_info.asp?" +
			"InNo=" + eventId + "&" +
			"MF=" + strconv.Itoa(MF)
		request = cookieAddedRequest(url, cookies)
		response, _ = client.Do(request)
		body, _ = ioutil.ReadAll(response.Body)
		html, error = iconv.Conv(string(body), "utf-8", "euc-kr")
		// 가끔 인코딩이 안될 때가 있음
		if error != nil {
			fmt.Println("인코딩에 실패했습니다. URL :", url)
			MF++
			eventId = nextEventId(string(body))
			if len(eventId) == 0 {
				break
			}
			continue
		}
		title := html[strings.Index(html, TITLE_START)+len(TITLE_START):]
		title = title[:strings.Index(title, TITLE_END)]
		content := ""
		tip := ""
		tipFind := false
		filtered := false
		company := ""
		announceMethod := ""
		var goods []string
		for _, entry := range filter {
			switch entry.typeValue {
			// 매일 응모
			case TYPE_DAILY:
				if strings.Index(html, DAILY) >= 0 {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 행사 내용
			case TYPE_CONTENT:
				if len(content) == 0 {
					content = html[strings.Index(html, CONTENT_START)+len(CONTENT_START):]
					content = strings.TrimSpace(content[:strings.Index(content, CONTENT_END)])
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				if regex.MatchString(content) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 응모 요령
			case TYPE_TIP:
				if !tipFind {
					index = strings.Index(html, TIP_START)
					if index >= 0 {
						tip = html[strings.Index(html, TIP_START)+len(TIP_START):]
						tip = strings.TrimSpace(tip[:strings.Index(tip, TIP_END)])
					}
					tipFind = true
				}
				if len(tip) > 0 {
					regex, error := regexp.Compile(entry.regex)
					if error != nil {
						fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
						return
					}
					if regex.MatchString(tip) {
						deleteEvent(eventId, cookies)
						filtered = true
					}
				}
			// 경품
			case TYPE_GOODS:
				if len(goods) == 0 {
					for goodsString := html; ; {
						index = strings.Index(goodsString, GOODS_START)
						if index < 0 {
							break
						}
						goodsString = goodsString[index+len(GOODS_START):]
						oneGoods := strings.TrimSpace(goodsString[:strings.Index(goodsString, GOODS_END)])
						goods = append(goods, oneGoods)
					}
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				matchCount := 0
				for _, oneGoods := range goods {
					if regex.MatchString(oneGoods) {
						matchCount++
					}
				}
				if matchCount == len(goods) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 주최사
			case TYPE_COMPANY:
				if len(company) == 0 {
					company = title[strings.Index(title, "[")+1 : strings.Index(title, "]")]
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				if regex.MatchString(company) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			// 발표방법
			case TYPE_ANNOUNCE_METHOD:
				if len(announceMethod) == 0 {
					announceMethod = html[strings.Index(html, ANNOUNCE_METHOD_START)+len(ANNOUNCE_METHOD_START):]
					announceMethod = strings.TrimSpace(announceMethod[:strings.Index(announceMethod, ANNOUNCE_METHOD_END)])
				}
				regex, error := regexp.Compile(entry.regex)
				if error != nil {
					fmt.Println("[필터 파일 오류]정규표현식을 해석하지 못했습니다. 해당 표현식 : " + entry.regex)
					return
				}
				if regex.MatchString(announceMethod) {
					deleteEvent(eventId, cookies)
					filtered = true
				}
			}
			if filtered {
				break
			}
		}
		fmt.Print(title + " => ")
		if filtered {
			fmt.Println("삭제")
		} else {
			fmt.Println("통과")
			MF++
		}
		eventId = nextEventId(html)
		if len(eventId) == 0 {
			break
		}
	}
}

func cookieAddedRequest(url string, cookies []*http.Cookie) *http.Request {
	request, _ := http.NewRequest("GET", url, nil)

	for _, cookie := range cookies {
		request.AddCookie(cookie)
	}

	return request
}

func deleteEvent(eventId string, cookies []*http.Cookie) {
	client := http.Client{}
	params := url.Values{"ApplyType": {"1,0"}, "InNo": {eventId}}
	request, _ := http.NewRequest("POST", "http://e-pass.co.kr/user/event_apply_detail.asp",
		bytes.NewBufferString(params.Encode()))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	for _, cookie := range cookies {
		request.AddCookie(cookie)
	}
	client.Do(request)
}

func nextEventId(html string) string {
	index := strings.Index(html, CURRENT_EVENT_ID_POSITION)
	if index < 0 {
		return ""
	}
	eventId := html[index:]
	eventId = eventId[strings.Index(eventId, EVENT_ID_START)+len(EVENT_ID_START):]
	eventId = eventId[strings.Index(eventId, EVENT_ID_START)+len(EVENT_ID_START):]
	index = strings.Index(eventId, EVENT_ID_START)
	if index < 0 {
		return ""
	}
	eventId = eventId[index+len(EVENT_ID_START):]
	eventId = eventId[:strings.Index(eventId, EVENT_ID_END)]

	return eventId
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/base64"
	"encoding/json"

	"github.com/abbot/go-http-auth"
	"github.com/dchest/uniuri"
	"github.com/zenazn/goji/bind"
	"goji.io"
	"goji.io/pat"
)

const usernameHeader string = "X-Authenticated-Username"
const cookieName string = "_htpasswd_auth"
const sessionFormat string = time.RFC3339
const slack time.Duration = time.Duration(10) * time.Second
const realm string = "example.com"

var sessionDir string
var domain string
var htpasswd string
var cookieLifetime int
var secure bool

var validChars map[rune]bool

func init() {
	validChars = map[rune]bool{}
	for _, c := range string(uniuri.StdChars) {
		validChars[c] = true
	}
}

type session struct {
	CreatedRaw string `json:"created"`
	Domain     string
	Username   string

	Name    string    `json:"-"`
	Created time.Time `json:"-"`
}

func (s *session) fill(name string) error {
	t, err := time.Parse(sessionFormat, s.CreatedRaw)
	if err != nil {
		return err
	}
	s.Created = t
	s.Name = name
	return nil
}

func (s *session) Valid() bool {
	sessionExpiry := s.Created.Add(time.Duration(cookieLifetime) * time.Second)
	return sessionExpiry.After(time.Now())
}
func validateCookie(cookie string) (*session, error) {
	for _, c := range cookie {
		if _, ok := validChars[c]; !ok {
			return nil, errors.New("Cookie value contains invalid characters.")
		}
	}
	sessionPath := filepath.Join(sessionDir, cookie)
	base, err := filepath.Rel(sessionDir, sessionPath)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(base, ".") || strings.HasPrefix(base, "/") {
		return nil, errors.New("Cookie value would traverse directories")
	}

	// From here on, the session pathname is assumed innocent & handle-able.
	data, err := ioutil.ReadFile(sessionPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Printf("Request for a session that doesn't exist: %q", sessionPath)
			return nil, err
		} else {
			log.Printf("Checking session %q failed: %s", sessionPath, err)
			return nil, err
		}
	}
	session := session{}
	json.Unmarshal(data, &session)
	err = session.fill(cookie)
	if err != nil {
		os.Remove(sessionPath) // clean up the session
		return nil, err
	}
	if !session.Valid() {
		return nil, errors.New("Session is no longer valid.")
	}
	return &session, nil
}

func invalidateCookie() *http.Cookie {
	return &http.Cookie{Name: cookieName, Value: "nope", Domain: domain, Path: "/", MaxAge: -1, Secure: secure, HttpOnly: true}
}

func checkSession(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false

	// Cleanup and ensure we always send a decent response:
	defer func() {
		if !success {
			http.SetCookie(w, newCookie)
			http.Error(w, "Session is not/no longer valid", http.StatusUnauthorized)
			return
		}
		w.WriteHeader(200)
		w.Write([]byte{})
	}()

	cookie, err := r.Cookie(cookieName)
	if err != nil {
		if err != http.ErrNoCookie {
			log.Printf("Failed getting the cookie: %s", err)
		}
		return

	}
	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		return
	}

	if session.Domain != r.Host {
		log.Printf("Session %q was made for another domain: %s; wanted %q", cookie.Value, session.Domain, r.Host)
		return
	}

	// We have a valid session that hasn't yet expired. Let's call it a success.
	w.Header().Set(usernameHeader, session.Username)
	success = true
	return
}

func newSession(domain, user string) (string, error) {
	name := uniuri.NewLen(90)
	sessionPath := filepath.Join(sessionDir, name)
	data, err := json.Marshal(session{
		CreatedRaw: time.Now().Add(slack).Format(sessionFormat),
		Domain:     domain,
		Username:   user,
	})
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(sessionPath, data, 0400)
	if err != nil {
		return "", err
	}
	return name, nil
}

func login(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false
	defer func() {
		http.SetCookie(w, newCookie)
		if !success {
			http.Error(w, "Nope", http.StatusForbidden)
			return
		}
	}()
	user := r.PostFormValue("login")
	if len(user) == 0 {
		return
	}

	password := r.PostFormValue("password")
	if len(password) == 0 {
		return
	}

	log.Printf("Validating login request by user %q on %q", user, r.Host)

	// This is very awful, but sadly go-http-auth has no better interface:
	authenticator := auth.NewBasicAuthenticator(realm, auth.HtpasswdFileProvider(htpasswd))
	r.Header["Authorization"] = []string{
		fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", user, password)))),
	}
	if user != authenticator.CheckAuth(r) {
		log.Printf("User %q couldn't be authenticated", user)
		return
	}

	// OK, we're authenticated. Let's get a session started:
	session, err := newSession(r.Host, user)
	if err != nil {
		return
	}
	newCookie.Value = session
	newCookie.MaxAge = cookieLifetime
	if r.PostFormValue("ephemeral") == "" {
		newCookie.Expires = time.Now().Add(time.Duration(cookieLifetime) * time.Second)
	}
	success = true
}

func runCleanup() {
	deleted := 0
	err := filepath.Walk(sessionDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			if path != sessionDir {
				return filepath.SkipDir
			} else {
				return nil
			}
		}
		data, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		var sess session
		err = json.Unmarshal(data, &sess)
		if err != nil {
			fmt.Printf("Session %q was invalid JSON (%s), leaving it", path, err)
			return nil
		}
		if err = sess.fill(filepath.Base(path)); err != nil {
			fmt.Printf("Session %q has an invalid date (%s), deleting it", path, err)
			return os.Remove(path)
		}
		if !sess.Valid() {
			deleted += 1
			return os.Remove(path)
		}
		return nil
	})
	if err != nil {
		log.Printf("Couldn't delete sessions: %s", err)
		os.Exit(2)
	}
	if deleted > 0 {
		log.Printf("Cleaned out %d sessions from %q.", deleted, sessionDir)
	}
}

func logout(w http.ResponseWriter, r *http.Request) {
	cookie, err := r.Cookie(cookieName)
	http.SetCookie(w, invalidateCookie())
	if err != nil {
		if err != http.ErrNoCookie {
			log.Printf("Failed getting the cookie: %s", err)
		}
		http.Redirect(w, r, "/", 301)
		return
	}

	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		http.Redirect(w, r, "/", 301)
		return
	}
	os.Remove(filepath.Join(sessionDir, session.Name))
	http.Redirect(w, r, "/", 301)
}

func main() {
	var cleanup bool

	bind.DefaultBind = "127.0.0.1:8000"
	bind.WithFlag()
	flag.StringVar(&sessionDir, "sessions", "/var/db/http-auth/cookies", "Directory in which htpasswd-login-form places sessions")
	flag.StringVar(&domain, "domain", "", "Domain to set on all cookies")
	flag.StringVar(&htpasswd, "htpasswd", "/etc/nginx/.htpasswd", "htpasswd file to use for authentication")
	flag.IntVar(&cookieLifetime, "lifetime", 86400, "Maximum cookie lifetime in seconds")
	flag.BoolVar(&secure, "secure", true, "Whether to set cookies to secure (false is useful for dev)")

	flag.BoolVar(&cleanup, "cleanup", false, "Perform once-in-a-while cleanup actions")
	flag.Parse()

	if cleanup {
		runCleanup()
		return
	}

	mux := goji.NewMux()
	mux.HandleFunc(pat.Get("/auth"), checkSession)
	mux.HandleFunc(pat.Post("/login"), login)
	mux.HandleFunc(pat.Post("/logout"), logout)

	http.Serve(bind.Default(), mux)
}
<commit_msg>Correctly set ephemeral cookies<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/base64"
	"encoding/json"

	"github.com/abbot/go-http-auth"
	"github.com/dchest/uniuri"
	"github.com/zenazn/goji/bind"
	"goji.io"
	"goji.io/pat"
)

const usernameHeader string = "X-Authenticated-Username"
const cookieName string = "_htpasswd_auth"
const sessionFormat string = time.RFC3339
const slack time.Duration = time.Duration(10) * time.Second
const realm string = "example.com"

var sessionDir string
var domain string
var htpasswd string
var cookieLifetime int
var secure bool

var validChars map[rune]bool

func init() {
	validChars = map[rune]bool{}
	for _, c := range string(uniuri.StdChars) {
		validChars[c] = true
	}
}

type session struct {
	CreatedRaw string `json:"created"`
	Domain     string
	Username   string

	Name    string    `json:"-"`
	Created time.Time `json:"-"`
}

func (s *session) fill(name string) error {
	t, err := time.Parse(sessionFormat, s.CreatedRaw)
	if err != nil {
		return err
	}
	s.Created = t
	s.Name = name
	return nil
}

func (s *session) Valid() bool {
	sessionExpiry := s.Created.Add(time.Duration(cookieLifetime) * time.Second)
	return sessionExpiry.After(time.Now())
}
func validateCookie(cookie string) (*session, error) {
	for _, c := range cookie {
		if _, ok := validChars[c]; !ok {
			return nil, errors.New("Cookie value contains invalid characters.")
		}
	}
	sessionPath := filepath.Join(sessionDir, cookie)
	base, err := filepath.Rel(sessionDir, sessionPath)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(base, ".") || strings.HasPrefix(base, "/") {
		return nil, errors.New("Cookie value would traverse directories")
	}

	// From here on, the session pathname is assumed innocent & handle-able.
	data, err := ioutil.ReadFile(sessionPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Printf("Request for a session that doesn't exist: %q", sessionPath)
			return nil, err
		} else {
			log.Printf("Checking session %q failed: %s", sessionPath, err)
			return nil, err
		}
	}
	session := session{}
	json.Unmarshal(data, &session)
	err = session.fill(cookie)
	if err != nil {
		os.Remove(sessionPath) // clean up the session
		return nil, err
	}
	if !session.Valid() {
		return nil, errors.New("Session is no longer valid.")
	}
	return &session, nil
}

func invalidateCookie() *http.Cookie {
	return &http.Cookie{Name: cookieName, Value: "nope", Domain: domain, Path: "/", MaxAge: -1, Secure: secure, HttpOnly: true}
}

func checkSession(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false

	// Cleanup and ensure we always send a decent response:
	defer func() {
		if !success {
			http.SetCookie(w, newCookie)
			http.Error(w, "Session is not/no longer valid", http.StatusUnauthorized)
			return
		}
		w.WriteHeader(200)
		w.Write([]byte{})
	}()

	cookie, err := r.Cookie(cookieName)
	if err != nil {
		if err != http.ErrNoCookie {
			log.Printf("Failed getting the cookie: %s", err)
		}
		return

	}
	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		return
	}

	if session.Domain != r.Host {
		log.Printf("Session %q was made for another domain: %s; wanted %q", cookie.Value, session.Domain, r.Host)
		return
	}

	// We have a valid session that hasn't yet expired. Let's call it a success.
	w.Header().Set(usernameHeader, session.Username)
	success = true
	return
}

func newSession(domain, user string) (string, error) {
	name := uniuri.NewLen(90)
	sessionPath := filepath.Join(sessionDir, name)
	data, err := json.Marshal(session{
		CreatedRaw: time.Now().Add(slack).Format(sessionFormat),
		Domain:     domain,
		Username:   user,
	})
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(sessionPath, data, 0400)
	if err != nil {
		return "", err
	}
	return name, nil
}

func login(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false
	defer func() {
		http.SetCookie(w, newCookie)
		if !success {
			http.Error(w, "Nope", http.StatusForbidden)
			return
		}
	}()
	user := r.PostFormValue("login")
	if len(user) == 0 {
		return
	}

	password := r.PostFormValue("password")
	if len(password) == 0 {
		return
	}

	log.Printf("Validating login request by user %q on %q", user, r.Host)

	// This is very awful, but sadly go-http-auth has no better interface:
	authenticator := auth.NewBasicAuthenticator(realm, auth.HtpasswdFileProvider(htpasswd))
	r.Header["Authorization"] = []string{
		fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", user, password)))),
	}
	if user != authenticator.CheckAuth(r) {
		log.Printf("User %q couldn't be authenticated", user)
		return
	}

	// OK, we're authenticated. Let's get a session started:
	session, err := newSession(r.Host, user)
	if err != nil {
		return
	}
	newCookie.Value = session
	if r.PostFormValue("ephemeral") == "" {
		newCookie.Expires = time.Now().Add(time.Duration(cookieLifetime) * time.Second)
		newCookie.MaxAge = cookieLifetime
	} else {
		// Ephemeral session requested - leave off any
		// MaxAge/Expires settings to attempt to make browsers
		// drop the cookie once the session ends:
		newCookie.MaxAge = 0
	}
	success = true
}

func runCleanup() {
	deleted := 0
	err := filepath.Walk(sessionDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			if path != sessionDir {
				return filepath.SkipDir
			} else {
				return nil
			}
		}
		data, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		var sess session
		err = json.Unmarshal(data, &sess)
		if err != nil {
			fmt.Printf("Session %q was invalid JSON (%s), leaving it", path, err)
			return nil
		}
		if err = sess.fill(filepath.Base(path)); err != nil {
			fmt.Printf("Session %q has an invalid date (%s), deleting it", path, err)
			return os.Remove(path)
		}
		if !sess.Valid() {
			deleted += 1
			return os.Remove(path)
		}
		return nil
	})
	if err != nil {
		log.Printf("Couldn't delete sessions: %s", err)
		os.Exit(2)
	}
	if deleted > 0 {
		log.Printf("Cleaned out %d sessions from %q.", deleted, sessionDir)
	}
}

func logout(w http.ResponseWriter, r *http.Request) {
	cookie, err := r.Cookie(cookieName)
	http.SetCookie(w, invalidateCookie())
	if err != nil {
		if err != http.ErrNoCookie {
			log.Printf("Failed getting the cookie: %s", err)
		}
		http.Redirect(w, r, "/", 301)
		return
	}

	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		http.Redirect(w, r, "/", 301)
		return
	}
	os.Remove(filepath.Join(sessionDir, session.Name))
	http.Redirect(w, r, "/", 301)
}

func main() {
	var cleanup bool

	bind.DefaultBind = "127.0.0.1:8000"
	bind.WithFlag()
	flag.StringVar(&sessionDir, "sessions", "/var/db/http-auth/cookies", "Directory in which htpasswd-login-form places sessions")
	flag.StringVar(&domain, "domain", "", "Domain to set on all cookies")
	flag.StringVar(&htpasswd, "htpasswd", "/etc/nginx/.htpasswd", "htpasswd file to use for authentication")
	flag.IntVar(&cookieLifetime, "lifetime", 86400, "Maximum cookie lifetime in seconds")
	flag.BoolVar(&secure, "secure", true, "Whether to set cookies to secure (false is useful for dev)")

	flag.BoolVar(&cleanup, "cleanup", false, "Perform once-in-a-while cleanup actions")
	flag.Parse()

	if cleanup {
		runCleanup()
		return
	}

	mux := goji.NewMux()
	mux.HandleFunc(pat.Get("/auth"), checkSession)
	mux.HandleFunc(pat.Post("/login"), login)
	mux.HandleFunc(pat.Post("/logout"), logout)

	http.Serve(bind.Default(), mux)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sort"
	"time"
)

import (
	"github.com/suntong/go-imap"
	"github.com/voxelbrain/goptions"
	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var (
	dumpProtocol *bool = flag.Bool("dumpprotocol", false, "dump imap stream")
)

////////////////////////////////////////////////////////////////////////////
// Function definitions

func check(err error) {
	if err != nil {
		panic(err)
	}
}

type progressMessage struct {
	cur, total int
	text       string
}

type UI struct {
	statusChan chan interface{}
	netmon     *netmonReader
}

func (ui *UI) log(format string, args ...interface{}) {
	ui.statusChan <- fmt.Sprintf(format, args...)
}

func (ui *UI) progress(cur, total int, format string, args ...interface{}) {
	message := &progressMessage{cur, total, fmt.Sprintf(format, args...)}
	ui.statusChan <- message
}

func loadAuth(path string) (string, string, string) {
	f, err := os.Open(path)
	check(err)
	r := bufio.NewReader(f)

	imap, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	user, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	pass, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	return string(imap), string(user), string(pass)
}

func readExtra(im *imap.IMAP) {
	for {
		select {
		case msg := <-im.Unsolicited:
			log.Printf("*** unsolicited: %T %+v", msg, msg)
		default:
			return
		}
	}
}

func (ui *UI) connect(useNetmon bool) *imap.IMAP {
	imapAddr, user, pass := loadAuth("auth")

	ui.log("connecting...")
	conn, err := tls.Dial("tcp", imapAddr, nil)
	check(err)

	var r io.Reader = conn
	if *dumpProtocol {
		r = newLoggingReader(r, 300)
	}
	if useNetmon {
		ui.netmon = newNetmonReader(r)
		r = ui.netmon
	}
	im := imap.New(r, conn)
	im.Unsolicited = make(chan interface{}, 100)

	hello, err := im.Start()
	check(err)
	ui.log("server hello: %s", hello)

	ui.log("logging in...")
	resp, caps, err := im.Auth(user, pass)
	check(err)
	ui.log("%s", resp)
	ui.log("server capabilities: %s", caps)

	return im
}

func (ui *UI) fetch(im *imap.IMAP, mailbox string) {
	ui.log("opening %s...", mailbox)
	examine, err := im.Examine(mailbox)
	check(err)
	ui.log("mailbox status: %+v", examine)
	readExtra(im)

	f, err := os.Create(mailbox + ".mbox")
	check(err)
	mbox := newMbox(f)

	query := fmt.Sprintf("1:%d", examine.Exists)
	ui.log("requesting messages %s", query)

	ch, err := im.FetchAsync(query, []string{"RFC822"})
	check(err)

	i := 0
	total := examine.Exists
	ui.progress(i, total, "fetching messages", i, total)
L:
	for {
		r := <-ch
		switch r := r.(type) {
		case *imap.ResponseFetch:
			if err = mbox.writeMessage(r.Rfc822, messageFilter); err != nil &&
				VERBOSITY > 1 {
				ui.log("message ignored: %v", err)
			}
			i++
			ui.progress(i, total, "fetching messages")
		case *imap.ResponseStatus:
			ui.log("complete %v\n", r)
			break L
		}
	}
	readExtra(im)
}

func (ui *UI) reportOnStatus() {

	ticker := time.NewTicker(1000 * 1000 * 1000)
	overprint := false
	status := ""
	overprintLast := false
	for ui.statusChan != nil {
		select {
		case s, stillOpen := <-ui.statusChan:
			switch s := s.(type) {
			case string:
				status = s
				overprint = false
			case *progressMessage:
				status = fmt.Sprintf("%s [%d/%d]", s.text, s.cur, s.total)
				overprint = true
			default:
				if s != nil {
					status = s.(error).Error()
					ui.statusChan = nil
					ticker.Stop()
				}
			}
			if !stillOpen {
				ui.statusChan = nil
				ticker.Stop()
			}
		case <-ticker.C:
			if ui.netmon != nil {
				ui.netmon.Tick()
			}
		}

		if overprintLast {
			fmt.Printf("\r\x1B[K")
		} else {
			fmt.Printf("\n")
		}
		overprintLast = overprint
		fmt.Printf("%s", status)
		if overprint && ui.netmon != nil {
			fmt.Printf(" [%.1fk/s]", ui.netmon.Bandwidth()/1000.0)
		}
	}
	fmt.Printf("\n")
}

func (ui *UI) runFetch(mailbox string) {
	ui.statusChan = make(chan interface{})
	go func() {
		defer func() {
			if e := recover(); e != nil {
				ui.statusChan <- e
			}
		}()
		im := ui.connect(true)
		ui.fetch(im, mailbox)
		close(ui.statusChan)
	}()

	ui.reportOnStatus()
}

func (ui *UI) runList() {
	ui.statusChan = make(chan interface{})
	go func() {
		im := ui.connect(false)
		mailboxes, err := im.List("", imap.WildcardAny)
		check(err)
		fmt.Printf("Available mailboxes:\n")
		for _, mailbox := range mailboxes {
			fmt.Printf("  %s\n", mailbox.Name)
		}
		readExtra(im)
	}()

	ui.reportOnStatus()
}

////////////////////////////////////////////////////////////////////////////
// Main

type Options struct {
	Verbosity []bool        `goptions:"-v, --verbose, description='Be verbose'"`
	Help      goptions.Help `goptions:"-h, --help, description='Show this help\n\nSub-commands (Verbs):\n\tlist\t\tList mailboxes\n\tfetch\t\tDownload mailbox'"`

	goptions.Verbs

	List struct{} `goptions:"list"`

	Fetch struct {
		Folder      string `goptions:"-f, --folder, description='Mail folder to fetch', obligatory"`
		TrackId     bool   `goptions:"-t, --trackid, description='Track message Id'"`
		WithinYear  int    `goptions:"--wy, description='Within years, only to fetch mails within this number of years'"`
		WithinMonth int    `goptions:"--wm, description='Within months, ditto for months'"`
		WithinDay   int    `goptions:"--wd, description='Within days, ditto for days'"`
	} `goptions:"fetch"`
}

var options = Options{ // Default values goes here
}

type Command func(Options) error

var commands = map[goptions.Verbs]Command{
	"list":  listCmd,
	"fetch": fetchCmd,
}

var (
	VERBOSITY = 0
)

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

	goptions.ParseAndFail(&options)
	//fmt.Printf("] %#v\r\n", options)

	if len(options.Verbs) == 0 {
		goptions.PrintHelp()
		os.Exit(2)
	}

	VERBOSITY = len(options.Verbosity)

	messageIds = make(map[string]bool)
	messageFetchMode = options.Verbs == "fetch"
	if cmd, found := commands[options.Verbs]; found {
		err := cmd(options)
		check(err)
	}

}

////////////////////////////////////////////////////////////////////////////
// Dispatch function definitions

func listCmd(options Options) error {
	ui := new(UI)
	ui.runList()
	println("done")
	return nil
}

func fetchCmd(options Options) error {
	ui := new(UI)

	if options.Fetch.WithinYear+options.Fetch.WithinMonth+
		options.Fetch.WithinDay > 0 {
		messageFilterFunc = validation(messageFilter)
		validFrom = time.Now()
		validFrom = validFrom.AddDate(-options.Fetch.WithinYear,
			-options.Fetch.WithinMonth, -options.Fetch.WithinDay)
	}

	ui.runFetch(options.Fetch.Folder)
	if options.Fetch.TrackId {
		msgIdSave()
	}
	return nil
}

////////////////////////////////////////////////////////////////////////////
// message filtering

type MsgIds struct {
	Msg []string
}

var (
	messageFilterFunc = validation(emptyFilter)
	messageIds        map[string]bool
	messageFetchMode  bool
	validFrom         time.Time
)

func messageFilter(rfc822 []byte, envelopeDate time.Time) error {
	err := filterEnvelopeDate(rfc822, envelopeDate)
	if err != nil {
		return err
	}
	return filterMsgId(rfc822, envelopeDate)
}

func filterEnvelopeDate(rfc822 []byte, envelopeDate time.Time) error {
	if envelopeDate.Before(validFrom) {
		return errors.New("older than picked date")
	}
	return nil
}

func filterMsgId(rfc822 []byte, envelopeDate time.Time) error {
	r := regexp.MustCompile(`(?i)\nMessage-ID: *<(.*?)> *\r*\n`).FindSubmatch(rfc822)
	if len(r) == 0 {
		panic("Internal error: Message-ID not found\n" + string(rfc822))
		os.Exit(1)
	}
	// fmt.Printf("\nr: %+v\n", r)
	msgId := string(r[1])
	if VERBOSITY > 0 {
		fmt.Printf("\nmsgId: %+v\n", msgId)
	}

	// if not in message fetching mode, check for existing MsgId first
	if !messageFetchMode && messageIds[msgId] {
		return errors.New("existing message Id")
	}

	messageIds[msgId] = true
	return nil
}

// msgIdSave will save out the global messageIds in sorted order
func msgIdSave() {
	// http://blog.golang.org/go-maps-in-action
	msgIds := MsgIds{}
	for k := range messageIds {
		msgIds.Msg = append(msgIds.Msg, k)
	}
	sort.Strings(msgIds.Msg)

	y, err := yaml.Marshal(&msgIds)
	check(err)

	// open output file
	f, err := os.Create(options.Fetch.Folder + ".yaml")
	check(err)
	// close fo on exit and check for its returned error
	defer func() {
		check(f.Close())
	}()

	fmt.Fprintf(f, "%s\n", string(y))
}
<commit_msg>- [+] add sync Sub-command having the same options as fetch<commit_after>package main

import (
	"bufio"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sort"
	"time"
)

import (
	"github.com/suntong/go-imap"
	"github.com/voxelbrain/goptions"
	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var (
	dumpProtocol *bool = flag.Bool("dumpprotocol", false, "dump imap stream")
)

////////////////////////////////////////////////////////////////////////////
// Function definitions

func check(err error) {
	if err != nil {
		panic(err)
	}
}

type progressMessage struct {
	cur, total int
	text       string
}

type UI struct {
	statusChan chan interface{}
	netmon     *netmonReader
}

func (ui *UI) log(format string, args ...interface{}) {
	ui.statusChan <- fmt.Sprintf(format, args...)
}

func (ui *UI) progress(cur, total int, format string, args ...interface{}) {
	message := &progressMessage{cur, total, fmt.Sprintf(format, args...)}
	ui.statusChan <- message
}

func loadAuth(path string) (string, string, string) {
	f, err := os.Open(path)
	check(err)
	r := bufio.NewReader(f)

	imap, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	user, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	pass, isPrefix, err := r.ReadLine()
	check(err)
	if isPrefix {
		panic("prefix")
	}

	return string(imap), string(user), string(pass)
}

func readExtra(im *imap.IMAP) {
	for {
		select {
		case msg := <-im.Unsolicited:
			log.Printf("*** unsolicited: %T %+v", msg, msg)
		default:
			return
		}
	}
}

func (ui *UI) connect(useNetmon bool) *imap.IMAP {
	imapAddr, user, pass := loadAuth("auth")

	ui.log("connecting...")
	conn, err := tls.Dial("tcp", imapAddr, nil)
	check(err)

	var r io.Reader = conn
	if *dumpProtocol {
		r = newLoggingReader(r, 300)
	}
	if useNetmon {
		ui.netmon = newNetmonReader(r)
		r = ui.netmon
	}
	im := imap.New(r, conn)
	im.Unsolicited = make(chan interface{}, 100)

	hello, err := im.Start()
	check(err)
	ui.log("server hello: %s", hello)

	ui.log("logging in...")
	resp, caps, err := im.Auth(user, pass)
	check(err)
	ui.log("%s", resp)
	ui.log("server capabilities: %s", caps)

	return im
}

func (ui *UI) fetch(im *imap.IMAP, mailbox string) {
	ui.log("opening %s...", mailbox)
	examine, err := im.Examine(mailbox)
	check(err)
	ui.log("mailbox status: %+v", examine)
	readExtra(im)

	f, err := os.Create(mailbox + ".mbox")
	check(err)
	mbox := newMbox(f)

	query := fmt.Sprintf("1:%d", examine.Exists)
	ui.log("requesting messages %s", query)

	ch, err := im.FetchAsync(query, []string{"RFC822"})
	check(err)

	i := 0
	total := examine.Exists
	ui.progress(i, total, "fetching messages", i, total)
L:
	for {
		r := <-ch
		switch r := r.(type) {
		case *imap.ResponseFetch:
			if err = mbox.writeMessage(r.Rfc822, messageFilter); err != nil &&
				VERBOSITY > 1 {
				ui.log("message ignored: %v", err)
			}
			i++
			ui.progress(i, total, "fetching messages")
		case *imap.ResponseStatus:
			ui.log("complete %v\n", r)
			break L
		}
	}
	readExtra(im)
}

func (ui *UI) reportOnStatus() {

	ticker := time.NewTicker(1000 * 1000 * 1000)
	overprint := false
	status := ""
	overprintLast := false
	for ui.statusChan != nil {
		select {
		case s, stillOpen := <-ui.statusChan:
			switch s := s.(type) {
			case string:
				status = s
				overprint = false
			case *progressMessage:
				status = fmt.Sprintf("%s [%d/%d]", s.text, s.cur, s.total)
				overprint = true
			default:
				if s != nil {
					status = s.(error).Error()
					ui.statusChan = nil
					ticker.Stop()
				}
			}
			if !stillOpen {
				ui.statusChan = nil
				ticker.Stop()
			}
		case <-ticker.C:
			if ui.netmon != nil {
				ui.netmon.Tick()
			}
		}

		if overprintLast {
			fmt.Printf("\r\x1B[K")
		} else {
			fmt.Printf("\n")
		}
		overprintLast = overprint
		fmt.Printf("%s", status)
		if overprint && ui.netmon != nil {
			fmt.Printf(" [%.1fk/s]", ui.netmon.Bandwidth()/1000.0)
		}
	}
	fmt.Printf("\n")
}

func (ui *UI) runFetch(mailbox string) {
	ui.statusChan = make(chan interface{})
	go func() {
		defer func() {
			if e := recover(); e != nil {
				ui.statusChan <- e
			}
		}()
		im := ui.connect(true)
		ui.fetch(im, mailbox)
		close(ui.statusChan)
	}()

	ui.reportOnStatus()
}

func (ui *UI) runList() {
	ui.statusChan = make(chan interface{})
	go func() {
		im := ui.connect(false)
		mailboxes, err := im.List("", imap.WildcardAny)
		check(err)
		fmt.Printf("Available mailboxes:\n")
		for _, mailbox := range mailboxes {
			fmt.Printf("  %s\n", mailbox.Name)
		}
		readExtra(im)
	}()

	ui.reportOnStatus()
}

////////////////////////////////////////////////////////////////////////////
// Main

type Fetch struct {
	Folder      string `goptions:"-f, --folder, description='Mail folder to fetch', obligatory"`
	TrackId     bool   `goptions:"-t, --trackid, description='Track message Id'"`
	WithinYear  int    `goptions:"--wy, description='Within years, only to fetch mails within this number of years'"`
	WithinMonth int    `goptions:"--wm, description='Within months, ditto for months'"`
	WithinDay   int    `goptions:"--wd, description='Within days, ditto for days'"`
}

type Options struct {
	Verbosity []bool        `goptions:"-v, --verbose, description='Be verbose'"`
	Help      goptions.Help `goptions:"-h, --help, description='Show this help\n\nSub-commands (Verbs):\n\tlist\t\tList mailboxes\n\tfetch\t\tDownload mailbox\n\tsync\t\tSync cloud mail folder to existing mailbox file'"`

	goptions.Verbs

	List struct{} `goptions:"list"`

	Fetch `goptions:"fetch"` // fields embedding
	Sync  Fetch              `goptions:"sync"`
}

var options = Options{ // Default values goes here
}

type Command func(Options) error

var commands = map[goptions.Verbs]Command{
	"list":  listCmd,
	"fetch": fetchCmd,
}

var (
	VERBOSITY = 0
)

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

	goptions.ParseAndFail(&options)
	//fmt.Printf("] %#v\r\n", options)

	if len(options.Verbs) == 0 {
		goptions.PrintHelp()
		os.Exit(2)
	}

	VERBOSITY = len(options.Verbosity)

	messageIds = make(map[string]bool)
	messageFetchMode = options.Verbs == "fetch"
	if cmd, found := commands[options.Verbs]; found {
		err := cmd(options)
		check(err)
	}

}

////////////////////////////////////////////////////////////////////////////
// Dispatch function definitions

func listCmd(options Options) error {
	ui := new(UI)
	ui.runList()
	println("done")
	return nil
}

func fetchCmd(options Options) error {
	ui := new(UI)

	if options.Fetch.WithinYear+options.Fetch.WithinMonth+
		options.Fetch.WithinDay > 0 {
		messageFilterFunc = validation(messageFilter)
		validFrom = time.Now()
		validFrom = validFrom.AddDate(-options.Fetch.WithinYear,
			-options.Fetch.WithinMonth, -options.Fetch.WithinDay)
	}

	ui.runFetch(options.Fetch.Folder)
	if options.Fetch.TrackId {
		msgIdSave()
	}
	return nil
}

////////////////////////////////////////////////////////////////////////////
// message filtering

type MsgIds struct {
	Msg []string
}

var (
	messageFilterFunc = validation(emptyFilter)
	messageIds        map[string]bool
	messageFetchMode  bool
	validFrom         time.Time
)

func messageFilter(rfc822 []byte, envelopeDate time.Time) error {
	err := filterEnvelopeDate(rfc822, envelopeDate)
	if err != nil {
		return err
	}
	return filterMsgId(rfc822, envelopeDate)
}

func filterEnvelopeDate(rfc822 []byte, envelopeDate time.Time) error {
	if envelopeDate.Before(validFrom) {
		return errors.New("older than picked date")
	}
	return nil
}

func filterMsgId(rfc822 []byte, envelopeDate time.Time) error {
	r := regexp.MustCompile(`(?i)\nMessage-ID: *<(.*?)> *\r*\n`).FindSubmatch(rfc822)
	if len(r) == 0 {
		panic("Internal error: Message-ID not found\n" + string(rfc822))
		os.Exit(1)
	}
	// fmt.Printf("\nr: %+v\n", r)
	msgId := string(r[1])
	if VERBOSITY > 0 {
		fmt.Printf("\nmsgId: %+v\n", msgId)
	}

	// if not in message fetching mode, check for existing MsgId first
	if !messageFetchMode && messageIds[msgId] {
		return errors.New("existing message Id")
	}

	messageIds[msgId] = true
	return nil
}

// msgIdSave will save out the global messageIds in sorted order
func msgIdSave() {
	// http://blog.golang.org/go-maps-in-action
	msgIds := MsgIds{}
	for k := range messageIds {
		msgIds.Msg = append(msgIds.Msg, k)
	}
	sort.Strings(msgIds.Msg)

	y, err := yaml.Marshal(&msgIds)
	check(err)

	// open output file
	f, err := os.Create(options.Fetch.Folder + ".yaml")
	check(err)
	// close fo on exit and check for its returned error
	defer func() {
		check(f.Close())
	}()

	fmt.Fprintf(f, "%s\n", string(y))
}
<|endoftext|>