<commit_before>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	WardenClient "github.com/cloudfoundry-incubator/garden/client"
	WardenConnection "github.com/cloudfoundry-incubator/garden/client/connection"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"

	"github.com/cloudfoundry-incubator/executor/maintain"
	steno "github.com/cloudfoundry/gosteno"
	"github.com/cloudfoundry/gunk/timeprovider"
	"github.com/cloudfoundry/storeadapter/etcdstoreadapter"
	"github.com/cloudfoundry/storeadapter/workerpool"

	"github.com/cloudfoundry-incubator/executor/configuration"
	"github.com/cloudfoundry-incubator/executor/executor"
	"github.com/cloudfoundry-incubator/executor/log_streamer_factory"
	Transformer "github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/executor/uploader"
	"github.com/pivotal-golang/archiver/compressor"
	"github.com/pivotal-golang/archiver/extractor"
	"github.com/pivotal-golang/cacheddownloader"
)

var executorAddress = flag.String(
	"executorAddress",
	"0.0.0.0:1700",
	"host:port to serve API requests on")

var containerOwnerName = flag.String(
	"containerOwnerName",
	"executor",
	"name to track containers created by this executor; they will be reaped on start",
)

var wardenNetwork = flag.String(
	"wardenNetwork",
	"unix",
	"network mode for warden server (tcp, unix)",
)

var wardenAddr = flag.String(
	"wardenAddr",
	"/tmp/warden.sock",
	"network address for warden server",
)

var etcdCluster = flag.String(
	"etcdCluster",
	"http://127.0.0.1:4001",
	"comma-separated list of etcd addresses (http://ip:port)",
)

var logLevel = flag.String(
	"logLevel",
	"info",
	"the logging level (none, fatal, error, warn, info, debug, debug1, debug2, all)",
)

var syslogName = flag.String(
	"syslogName",
	"",
	"syslog name",
)

var memoryMBFlag = flag.String(
	"memoryMB",
	configuration.Automatic,
	"the amount of memory the executor has available in megabytes",
)

var diskMBFlag = flag.String(
	"diskMB",
	configuration.Automatic,
	"the amount of disk the executor has available in megabytes",
)

var heartbeatInterval = flag.Duration(
	"heartbeatInterval",
	60*time.Second,
	"the interval, in seconds, between heartbeats for maintaining presence",
)

var tempDir = flag.String(
	"tempDir",
	"/tmp",
	"location to store temporary assets",
)

var loggregatorServer = flag.String(
	"loggregatorServer",
	"",
	"loggregator server to emit logs to",
)

var loggregatorSecret = flag.String(
	"loggregatorSecret",
	"",
	"secret for the loggregator server",
)

var stack = flag.String(
	"stack",
	"",
	"the executor stack - must be specified",
)

var drainTimeout = flag.Duration(
	"drainTimeout",
	15*time.Minute,
	"time to give running tasks to drain before exiting",
)

var containerInodeLimit = flag.Int(
	"containerInodeLimit",
	200000,
	"max number of inodes per container",
)

var containerMaxCpuShares = flag.Int(
	"containerMaxCpuShares",
	0,
	"cpu shares allocatable to a container",
)

var cachePath = flag.String(
	"cachePath",
	"/tmp/cache",
	"location to cache assets",
)

var maxCacheSizeInBytes = flag.Int64(
	"maxCacheSizeInBytes",
	10*1024*1024*1024,
	"maximum size of the cache (in bytes) - you should include a healthy amount of overhead",
)

func main() {
	flag.Parse()

	if *stack == "" {
		log.Fatalf("A stack must be specified")
	}

	l, err := steno.GetLogLevel(*logLevel)
	if err != nil {
		log.Fatalf("Invalid loglevel: %s\n", *logLevel)
	}

	stenoConfig := steno.Config{
		Level: l,
		Sinks: []steno.Sink{steno.NewIOSink(os.Stdout)},
	}

	if *syslogName != "" {
		stenoConfig.Sinks = append(stenoConfig.Sinks, steno.NewSyslogSink(*syslogName))
	}

	steno.Init(&stenoConfig)
	logger := steno.NewLogger("executor")

	etcdAdapter := etcdstoreadapter.NewETCDStoreAdapter(
		strings.Split(*etcdCluster, ","),
		workerpool.NewWorkerPool(10),
	)

	bbs := Bbs.New(etcdAdapter, timeprovider.NewTimeProvider())

	err = etcdAdapter.Connect()
	if err != nil {
		logger.Errord(map[string]interface{}{
			"error": err,
		}, "failed to get etcdAdapter to connect")
		os.Exit(1)
	}

	wardenClient := WardenClient.New(&WardenConnection.Info{
		Network: *wardenNetwork,
		Addr:    *wardenAddr,
	})

	config := configuration.New(wardenClient)
	memoryMB, err := config.GetMemoryInMB(*memoryMBFlag)
	if err != nil {
		logger.Error(err.Error())
		os.Exit(1)
	}
	logger.Infof("Using memory limit %dMB", memoryMB)

	diskMB, err := config.GetDiskInMB(*diskMBFlag)
	if err != nil {
		logger.Error(err.Error())
		os.Exit(1)
	}
	logger.Infof("Using disk limit %dMB", diskMB)

	if *containerMaxCpuShares <= 0 {
		logger.Error("valid maximum container cpu shares must be specified on startup!")
		os.Exit(1)
	}

	cache := cacheddownloader.New(*cachePath, *tempDir, *maxCacheSizeInBytes, 10*time.Minute)
	uploader := uploader.New(10*time.Minute, logger)
	extractor := extractor.NewDetectable()
	compressor := compressor.NewTgz()
	logStreamerFactory := log_streamer_factory.New(
		*loggregatorServer,
		*loggregatorSecret,
	)

	transformer := Transformer.NewTransformer(
		logStreamerFactory,
		cache,
		uploader,
		extractor,
		compressor,
		logger,
		*tempDir,
	)

	logger.Info("executor.starting")

	executor := executor.New(
		*executorAddress,
		*containerOwnerName,
		uint64(*containerMaxCpuShares),
		memoryMB,
		diskMB,
		wardenClient,
		transformer,
		*drainTimeout,
		logger,
	)

	maintainSignals := make(chan os.Signal, 1)
	signal.Notify(maintainSignals, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)

	executorSignals := make(chan os.Signal, 1)
	signal.Notify(executorSignals, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)

	maintainReady := make(chan struct{})
	executorReady := make(chan struct{})

	maintainer := maintain.New(executor.ID(), bbs, logger, *heartbeatInterval)

	go func() {
		err := maintainer.Run(maintainSignals, maintainReady)
		if err != nil {
			logger.Errorf("failed to start maintaining presence: %s", err.Error())
			executorSignals <- syscall.SIGTERM
		}
	}()

	<-maintainReady

	go func() {
		<-executorReady
		logger.Info("executor.started")
	}()

	err = executor.Run(executorSignals, executorReady)

	if err != nil {
		os.Exit(1)
	}
	os.Exit(0)
}
<commit_msg>executorAddress -> listenAddr<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	WardenClient "github.com/cloudfoundry-incubator/garden/client"
	WardenConnection "github.com/cloudfoundry-incubator/garden/client/connection"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"

	"github.com/cloudfoundry-incubator/executor/maintain"
	steno "github.com/cloudfoundry/gosteno"
	"github.com/cloudfoundry/gunk/timeprovider"
	"github.com/cloudfoundry/storeadapter/etcdstoreadapter"
	"github.com/cloudfoundry/storeadapter/workerpool"

	"github.com/cloudfoundry-incubator/executor/configuration"
	"github.com/cloudfoundry-incubator/executor/executor"
	"github.com/cloudfoundry-incubator/executor/log_streamer_factory"
	Transformer "github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/executor/uploader"
	"github.com/pivotal-golang/archiver/compressor"
	"github.com/pivotal-golang/archiver/extractor"
	"github.com/pivotal-golang/cacheddownloader"
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:1700",
	"host:port to serve API requests on")

var containerOwnerName = flag.String(
	"containerOwnerName",
	"executor",
	"name to track containers created by this executor; they will be reaped on start",
)

var wardenNetwork = flag.String(
	"wardenNetwork",
	"unix",
	"network mode for warden server (tcp, unix)",
)

var wardenAddr = flag.String(
	"wardenAddr",
	"/tmp/warden.sock",
	"network address for warden server",
)

var etcdCluster = flag.String(
	"etcdCluster",
	"http://127.0.0.1:4001",
	"comma-separated list of etcd addresses (http://ip:port)",
)

var logLevel = flag.String(
	"logLevel",
	"info",
	"the logging level (none, fatal, error, warn, info, debug, debug1, debug2, all)",
)

var syslogName = flag.String(
	"syslogName",
	"",
	"syslog name",
)

var memoryMBFlag = flag.String(
	"memoryMB",
	configuration.Automatic,
	"the amount of memory the executor has available in megabytes",
)

var diskMBFlag = flag.String(
	"diskMB",
	configuration.Automatic,
	"the amount of disk the executor has available in megabytes",
)

var heartbeatInterval = flag.Duration(
	"heartbeatInterval",
	60*time.Second,
	"the interval, in seconds, between heartbeats for maintaining presence",
)

var tempDir = flag.String(
	"tempDir",
	"/tmp",
	"location to store temporary assets",
)

var loggregatorServer = flag.String(
	"loggregatorServer",
	"",
	"loggregator server to emit logs to",
)

var loggregatorSecret = flag.String(
	"loggregatorSecret",
	"",
	"secret for the loggregator server",
)

var stack = flag.String(
	"stack",
	"",
	"the executor stack - must be specified",
)

var drainTimeout = flag.Duration(
	"drainTimeout",
	15*time.Minute,
	"time to give running tasks to drain before exiting",
)

var containerInodeLimit = flag.Int(
	"containerInodeLimit",
	200000,
	"max number of inodes per container",
)

var containerMaxCpuShares = flag.Int(
	"containerMaxCpuShares",
	0,
	"cpu shares allocatable to a container",
)

var cachePath = flag.String(
	"cachePath",
	"/tmp/cache",
	"location to cache assets",
)

var maxCacheSizeInBytes = flag.Int64(
	"maxCacheSizeInBytes",
	10*1024*1024*1024,
	"maximum size of the cache (in bytes) - you should include a healthy amount of overhead",
)

func main() {
	flag.Parse()

	if *stack == "" {
		log.Fatalf("A stack must be specified")
	}

	l, err := steno.GetLogLevel(*logLevel)
	if err != nil {
		log.Fatalf("Invalid loglevel: %s\n", *logLevel)
	}

	stenoConfig := steno.Config{
		Level: l,
		Sinks: []steno.Sink{steno.NewIOSink(os.Stdout)},
	}

	if *syslogName != "" {
		stenoConfig.Sinks = append(stenoConfig.Sinks, steno.NewSyslogSink(*syslogName))
	}

	steno.Init(&stenoConfig)
	logger := steno.NewLogger("executor")

	etcdAdapter := etcdstoreadapter.NewETCDStoreAdapter(
		strings.Split(*etcdCluster, ","),
		workerpool.NewWorkerPool(10),
	)

	bbs := Bbs.New(etcdAdapter, timeprovider.NewTimeProvider())

	err = etcdAdapter.Connect()
	if err != nil {
		logger.Errord(map[string]interface{}{
			"error": err,
		}, "failed to get etcdAdapter to connect")
		os.Exit(1)
	}

	wardenClient := WardenClient.New(&WardenConnection.Info{
		Network: *wardenNetwork,
		Addr:    *wardenAddr,
	})

	config := configuration.New(wardenClient)
	memoryMB, err := config.GetMemoryInMB(*memoryMBFlag)
	if err != nil {
		logger.Error(err.Error())
		os.Exit(1)
	}
	logger.Infof("Using memory limit %dMB", memoryMB)

	diskMB, err := config.GetDiskInMB(*diskMBFlag)
	if err != nil {
		logger.Error(err.Error())
		os.Exit(1)
	}
	logger.Infof("Using disk limit %dMB", diskMB)

	if *containerMaxCpuShares <= 0 {
		logger.Error("valid maximum container cpu shares must be specified on startup!")
		os.Exit(1)
	}

	cache := cacheddownloader.New(*cachePath, *tempDir, *maxCacheSizeInBytes, 10*time.Minute)
	uploader := uploader.New(10*time.Minute, logger)
	extractor := extractor.NewDetectable()
	compressor := compressor.NewTgz()
	logStreamerFactory := log_streamer_factory.New(
		*loggregatorServer,
		*loggregatorSecret,
	)

	transformer := Transformer.NewTransformer(
		logStreamerFactory,
		cache,
		uploader,
		extractor,
		compressor,
		logger,
		*tempDir,
	)

	logger.Info("executor.starting")

	executor := executor.New(
		*listenAddr,
		*containerOwnerName,
		uint64(*containerMaxCpuShares),
		memoryMB,
		diskMB,
		wardenClient,
		transformer,
		*drainTimeout,
		logger,
	)

	maintainSignals := make(chan os.Signal, 1)
	signal.Notify(maintainSignals, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)

	executorSignals := make(chan os.Signal, 1)
	signal.Notify(executorSignals, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)

	maintainReady := make(chan struct{})
	executorReady := make(chan struct{})

	maintainer := maintain.New(executor.ID(), bbs, logger, *heartbeatInterval)

	go func() {
		err := maintainer.Run(maintainSignals, maintainReady)
		if err != nil {
			logger.Errorf("failed to start maintaining presence: %s", err.Error())
			executorSignals <- syscall.SIGTERM
		}
	}()

	<-maintainReady

	go func() {
		<-executorReady
		logger.Info("executor.started")
	}()

	err = executor.Run(executorSignals, executorReady)

	if err != nil {
		os.Exit(1)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/freeusd/solebtc/Godeps/_workspace/src/github.com/gin-gonic/gin"
	"github.com/freeusd/solebtc/Godeps/_workspace/src/github.com/robfig/cron"
	"github.com/freeusd/solebtc/controllers/v1"
	"github.com/freeusd/solebtc/errors"
	"github.com/freeusd/solebtc/middlewares"
	"github.com/freeusd/solebtc/models"
	"github.com/freeusd/solebtc/services/cache"
	"github.com/freeusd/solebtc/services/cache/memory"
	"github.com/freeusd/solebtc/services/mail"
	"github.com/freeusd/solebtc/services/mail/mandrill"
	"github.com/freeusd/solebtc/services/storage"
	"github.com/freeusd/solebtc/services/storage/mysql"
	"github.com/freeusd/solebtc/utils"
)

var (
	logWriter   io.Writer = os.Stdout
	panicWriter io.Writer = os.Stderr
	mailer      mail.Mailer
	store       storage.Storage
	memoryCache cache.Cache
)

func init() {
	initConfig()
	initMailer()
	initStorage()
	initCache()
	initCronjob()
}

func main() {
	gin.SetMode(ginEnvMode())
	router := gin.New()

	// middlewares
	recovery := gin.RecoveryWithWriter(panicWriter)
	logger := middlewares.LoggerWithWriter(logWriter)
	cors := middlewares.CORS()
	errorWriter := middlewares.ErrorWriter()
	authRequired := middlewares.AuthRequired(store.GetAuthToken, config.AuthToken.Lifetime)

	// globally use middlewares
	router.Use(recovery, logger, cors, errorWriter)

	// version 1 api endpoints
	v1Endpoints := router.Group("/v1")

	// user endpoints
	v1UserEndpoints := v1Endpoints.Group("/users")
	v1UserEndpoints.GET("", authRequired, v1.UserInfo(store.GetUserByID))
	v1UserEndpoints.POST("", v1.Signup(store.CreateUser, store.GetUserByID))
	v1UserEndpoints.PUT("/:id/status", v1.VerifyEmail(store.GetSessionByToken, store.GetUserByID, store.UpdateUserStatus))
	v1UserEndpoints.GET("/referees", authRequired, v1.RefereeList(store.GetRefereesSince, store.GetRefereesUntil))

	// auth token endpoints
	v1AuthTokenEndpoints := v1Endpoints.Group("/auth_tokens")
	v1AuthTokenEndpoints.POST("", v1.Login(store.GetUserByEmail, store.CreateAuthToken))
	v1AuthTokenEndpoints.DELETE("", authRequired, v1.Logout(store.DeleteAuthToken))

	// session endpoints
	v1SessionEndpoints := v1Endpoints.Group("/sessions")
	v1SessionEndpoints.POST("", authRequired, v1.RequestVerifyEmail(store.GetUserByID, store.UpsertSession, mailer.SendEmail))

	// income endpoints
	v1IncomeEndpoints := v1Endpoints.Group("/incomes", authRequired)
	v1IncomeEndpoints.POST("/rewards",
		v1.GetReward(store.GetUserByID,
			memoryCache.GetLatestTotalReward,
			memoryCache.GetLatestConfig,
			memoryCache.GetRewardRatesByType,
			createRewardIncome))
	v1IncomeEndpoints.GET("/rewards", v1.RewardList(store.GetRewardIncomesSince, store.GetRewardIncomesUntil))
	v1IncomeEndpoints.GET("/rewards/referees/:referee_id", v1.RefereeRewardList(store.GetUserByID, store.GetRewardIncomesSince, store.GetRewardIncomesUntil))

	fmt.Fprintf(logWriter, "SoleBTC is running on %s\n", config.HTTP.Port)
	panicIfErrored(router.Run(config.HTTP.Port))
}

func createRewardIncome(income models.Income, now time.Time) *errors.Error {
	if err := store.CreateRewardIncome(income, now); err != nil {
		return err
	}

	totalReward := income.Income
	if income.RefererID > 0 {
		totalReward += income.RefererIncome
	}
	memoryCache.IncrementTotalReward(now, totalReward)

	return nil
}

func initMailer() {
	// mailer
	mailer = mandrill.New(config.Mandrill.Key, config.Mandrill.FromEmail, config.Mandrill.FromName)
}

func initStorage() {
	// storage service
	var err error
	store, err = mysql.New(config.DB.DataSourceName)
	panicIfErrored(err)
}

func initCache() {
	memoryCache = memory.New()

	// init config in cache
	config, err := store.GetLatestConfig()
	panicIfErrored(err)
	memoryCache.SetLatestConfig(config)

	// init rates in cache
	lessRates, err := store.GetRewardRatesByType(models.RewardRateTypeLess)
	panicIfErrored(err)
	memoryCache.SetRewardRates(models.RewardRateTypeLess, lessRates)

	moreRates, err := store.GetRewardRatesByType(models.RewardRateTypeMore)
	panicIfErrored(err)
	memoryCache.SetRewardRates(models.RewardRateTypeMore, moreRates)

	// update bitcoin price on start
	updateBitcoinPrice()
}

func initCronjob() {
	c := cron.New()
	panicIfErrored(c.AddFunc("@every 1m", updateBitcoinPrice))
	c.Start()
}

// update bitcoin price in cache
func updateBitcoinPrice() {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintf(logWriter, "Update bitcoin price panic: %v\n", err)
		}
	}()

	// get bitcoin price from blockchain.info
	p, err := utils.BitcoinPrice()
	if err != nil {
		fmt.Fprintf(logWriter, "Fetch bitcoin price error: %v\n", err)
		return
	}

	// update bitcoin price in database
	if err := store.UpdateLatestBitcoinPrice(p); err != nil {
		fmt.Fprintf(logWriter, "Update bitcoin price in database error: %v\n", err)
		return
	}

	// update bitcoin price in cache
	memoryCache.UpdateBitcoinPrice(p)

	fmt.Fprintf(logWriter, "Successfully update bitcoin price to %v\n", p)
}

// fail fast on initialization
func panicIfErrored(err error) {
	if err != nil {
		panic(err)
	}
}
<commit_msg>Run create withdrawls everyday on midnight<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/freeusd/solebtc/Godeps/_workspace/src/github.com/gin-gonic/gin"
	"github.com/freeusd/solebtc/Godeps/_workspace/src/github.com/robfig/cron"
	"github.com/freeusd/solebtc/controllers/v1"
	"github.com/freeusd/solebtc/errors"
	"github.com/freeusd/solebtc/middlewares"
	"github.com/freeusd/solebtc/models"
	"github.com/freeusd/solebtc/services/cache"
	"github.com/freeusd/solebtc/services/cache/memory"
	"github.com/freeusd/solebtc/services/mail"
	"github.com/freeusd/solebtc/services/mail/mandrill"
	"github.com/freeusd/solebtc/services/storage"
	"github.com/freeusd/solebtc/services/storage/mysql"
	"github.com/freeusd/solebtc/utils"
)

var (
	logWriter   io.Writer = os.Stdout
	panicWriter io.Writer = os.Stderr
	mailer      mail.Mailer
	store       storage.Storage
	memoryCache cache.Cache
)

func init() {
	// ORDER MATTERs
	initConfig()
	initMailer()
	initStorage()
	initCache()
	initCronjob()
}

func main() {
	gin.SetMode(ginEnvMode())
	router := gin.New()

	// middlewares
	recovery := gin.RecoveryWithWriter(panicWriter)
	logger := middlewares.LoggerWithWriter(logWriter)
	cors := middlewares.CORS()
	errorWriter := middlewares.ErrorWriter()
	authRequired := middlewares.AuthRequired(store.GetAuthToken, config.AuthToken.Lifetime)

	// globally use middlewares
	router.Use(recovery, logger, cors, errorWriter)

	// version 1 api endpoints
	v1Endpoints := router.Group("/v1")

	// user endpoints
	v1UserEndpoints := v1Endpoints.Group("/users")
	v1UserEndpoints.GET("", authRequired, v1.UserInfo(store.GetUserByID))
	v1UserEndpoints.POST("", v1.Signup(store.CreateUser, store.GetUserByID))
	v1UserEndpoints.PUT("/:id/status", v1.VerifyEmail(store.GetSessionByToken, store.GetUserByID, store.UpdateUserStatus))
	v1UserEndpoints.GET("/referees", authRequired, v1.RefereeList(store.GetRefereesSince, store.GetRefereesUntil))

	// auth token endpoints
	v1AuthTokenEndpoints := v1Endpoints.Group("/auth_tokens")
	v1AuthTokenEndpoints.POST("", v1.Login(store.GetUserByEmail, store.CreateAuthToken))
	v1AuthTokenEndpoints.DELETE("", authRequired, v1.Logout(store.DeleteAuthToken))

	// session endpoints
	v1SessionEndpoints := v1Endpoints.Group("/sessions")
	v1SessionEndpoints.POST("", authRequired, v1.RequestVerifyEmail(store.GetUserByID, store.UpsertSession, mailer.SendEmail))

	// income endpoints
	v1IncomeEndpoints := v1Endpoints.Group("/incomes", authRequired)
	v1IncomeEndpoints.POST("/rewards",
		v1.GetReward(store.GetUserByID,
			memoryCache.GetLatestTotalReward,
			memoryCache.GetLatestConfig,
			memoryCache.GetRewardRatesByType,
			createRewardIncome))
	v1IncomeEndpoints.GET("/rewards", v1.RewardList(store.GetRewardIncomesSince, store.GetRewardIncomesUntil))
	v1IncomeEndpoints.GET("/rewards/referees/:referee_id", v1.RefereeRewardList(store.GetUserByID, store.GetRewardIncomesSince, store.GetRewardIncomesUntil))

	fmt.Fprintf(logWriter, "SoleBTC is running on %s\n", config.HTTP.Port)
	panicIfErrored(router.Run(config.HTTP.Port))
}

func createRewardIncome(income models.Income, now time.Time) *errors.Error {
	if err := store.CreateRewardIncome(income, now); err != nil {
		return err
	}

	totalReward := income.Income
	if income.RefererID > 0 {
		totalReward += income.RefererIncome
	}
	memoryCache.IncrementTotalReward(now, totalReward)

	return nil
}

func initMailer() {
	// mailer
	mailer = mandrill.New(config.Mandrill.Key, config.Mandrill.FromEmail, config.Mandrill.FromName)
}

func initStorage() {
	// storage service
	var err error
	store, err = mysql.New(config.DB.DataSourceName)
	panicIfErrored(err)
}

func initCache() {
	memoryCache = memory.New()

	// init config in cache
	config, err := store.GetLatestConfig()
	panicIfErrored(err)
	memoryCache.SetLatestConfig(config)

	// init rates in cache
	lessRates, err := store.GetRewardRatesByType(models.RewardRateTypeLess)
	panicIfErrored(err)
	memoryCache.SetRewardRates(models.RewardRateTypeLess, lessRates)

	moreRates, err := store.GetRewardRatesByType(models.RewardRateTypeMore)
	panicIfErrored(err)
	memoryCache.SetRewardRates(models.RewardRateTypeMore, moreRates)

	// update bitcoin price on start
	updateBitcoinPrice()
}

func initCronjob() {
	c := cron.New()
	panicIfErrored(c.AddFunc("@every 1m", updateBitcoinPrice))
	panicIfErrored(c.AddFunc("@daily", createWithdrawl))
	c.Start()
}

// update bitcoin price in cache
func updateBitcoinPrice() {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintf(logWriter, "Update bitcoin price panic: %v\n", err)
		}
	}()

	// get bitcoin price from blockchain.info
	p, err := utils.BitcoinPrice()
	if err != nil {
		fmt.Fprintf(logWriter, "Fetch bitcoin price error: %v\n", err)
		return
	}

	// update bitcoin price in database
	if err := store.UpdateLatestBitcoinPrice(p); err != nil {
		fmt.Fprintf(logWriter, "Update bitcoin price in database error: %v\n", err)
		return
	}

	// update bitcoin price in cache
	memoryCache.UpdateBitcoinPrice(p)

	fmt.Fprintf(logWriter, "Successfully update bitcoin price to %v\n", p)
}

// automatically create withdrawl
func createWithdrawl() {
	// critical for the system
	// use panicWriter for error log instead of logWriter
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintf(panicWriter, "Create withdrawl panic: %v\n", err)
		}
	}()

	users, err := store.GetWithdrawableUsers()
	if err != nil {
		fmt.Fprintf(panicWriter, "Get withdrawable users error: %v\n", err)
		return
	}

	f := func(users []models.User, handler func(err error, u models.User)) {
		for i := range users {
			handler(store.CreateWithdrawl(models.Withdrawl{
				UserID:         users[i].ID,
				Amount:         users[i].Balance,
				BitcoinAddress: users[i].BitcoinAddress,
			}), users[i])
		}
	}

	// create withdrawl, move errored ones into retry queue
	retryUsers := []models.User{}
	f(users, func(err error, u models.User) {
		if err != nil {
			retryUsers = append(retryUsers, u)
		}
	})

	// retry with error output
	errored := false
	f(retryUsers, func(err error, u models.User) {
		if err != nil {
			fmt.Fprintf(panicWriter, "Create withdrawl for user %v error: %v\n", u, err)
			errored = true
		}
	})

	if !errored {
		fmt.Fprintf(logWriter, "Create withdrawls successfully...\n")
	}
}

// fail fast on initialization
func panicIfErrored(err error) {
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"flag"
	"fmt"
	"github.com/jaksi/sshesame/channel"
	"github.com/jaksi/sshesame/request"
	"golang.org/x/crypto/ed25519"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"log"
	"net"
)

func main() {
	hostKey := flag.String("host_key", "", "a file containing a private key to use")
	listenAddress := flag.String("listen_address", "localhost", "the local address to listen on")
	port := flag.Uint("port", 2022, "the port number to listen on")
	flag.Parse()

	var key ssh.Signer
	var err error
	if *hostKey != "" {
		keyBytes, err := ioutil.ReadFile(*hostKey)
		if err != nil {
			log.Fatalln("Failed to read host key:", err.Error())
		}
		key, err = ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			log.Fatalln("Failed to parse host key:", err.Error())
		}
	} else {
		log.Println("WARNING: Generating a temporary private key. Consider creating one and passing it to -host_key")
		_, keyBytes, err := ed25519.GenerateKey(nil)
		if err != nil {
			log.Fatalln("Failed to generate temporary private key:", err.Error())
		}
		key, err = ssh.NewSignerFromSigner(keyBytes)
		if err != nil {
			log.Fatalln("Failed to parse generated private key:", err.Error())
		}
		log.Printf("SHA-256 fingerprint: %v\n", sha256.Sum256(key.PublicKey().Marshal()))
	}

	serverConfig := &ssh.ServerConfig{
		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			log.Printf("Login: client=%v, user=%q, password=%q\n", conn.RemoteAddr(), conn.User(), password)
			return nil, nil
		},
	}
	serverConfig.AddHostKey(key)

	listener, err := net.Listen("tcp", fmt.Sprintf("%v:%v", *listenAddress, *port))
	if err != nil {
		log.Fatalln("Failed to listen:", err.Error())
	}
	log.Printf("Listen: %v\n", listener.Addr())
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept connection:", err.Error())
			continue
		}
		log.Printf("Connection: client=%v\n", conn.RemoteAddr())
		go handleConn(serverConfig, conn)
	}
}

func handleConn(serverConfig *ssh.ServerConfig, conn net.Conn) {
	defer conn.Close()
	_, channels, requests, err := ssh.NewServerConn(conn, serverConfig)
	if err != nil {
		log.Println("Failed to establish SSH connection:", err.Error())
		return
	}
	log.Printf("Established SSH connection: client=%v\n", conn.RemoteAddr())
	go request.Handle(conn.RemoteAddr(), "global", requests)
	for newChannel := range channels {
		go channel.Handle(conn.RemoteAddr(), newChannel)
	}
	log.Printf("Disconnect: client=%v\n", conn.RemoteAddr())
}
<commit_msg>Add the client version to the password callback log<commit_after>package main

import (
	"crypto/sha256"
	"flag"
	"fmt"
	"github.com/jaksi/sshesame/channel"
	"github.com/jaksi/sshesame/request"
	"golang.org/x/crypto/ed25519"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"log"
	"net"
)

func main() {
	hostKey := flag.String("host_key", "", "a file containing a private key to use")
	listenAddress := flag.String("listen_address", "localhost", "the local address to listen on")
	port := flag.Uint("port", 2022, "the port number to listen on")
	flag.Parse()

	var key ssh.Signer
	var err error
	if *hostKey != "" {
		keyBytes, err := ioutil.ReadFile(*hostKey)
		if err != nil {
			log.Fatalln("Failed to read host key:", err.Error())
		}
		key, err = ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			log.Fatalln("Failed to parse host key:", err.Error())
		}
	} else {
		log.Println("WARNING: Generating a temporary private key. Consider creating one and passing it to -host_key")
		_, keyBytes, err := ed25519.GenerateKey(nil)
		if err != nil {
			log.Fatalln("Failed to generate temporary private key:", err.Error())
		}
		key, err = ssh.NewSignerFromSigner(keyBytes)
		if err != nil {
			log.Fatalln("Failed to parse generated private key:", err.Error())
		}
		log.Printf("SHA-256 fingerprint: %v\n", sha256.Sum256(key.PublicKey().Marshal()))
	}

	serverConfig := &ssh.ServerConfig{
		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			log.Printf("Login: client=%v, user=%q, password=%q, version=%q\n", conn.RemoteAddr(), conn.User(), password, conn.ClientVersion())
			return nil, nil
		},
	}
	serverConfig.AddHostKey(key)

	listener, err := net.Listen("tcp", fmt.Sprintf("%v:%v", *listenAddress, *port))
	if err != nil {
		log.Fatalln("Failed to listen:", err.Error())
	}
	log.Printf("Listen: %v\n", listener.Addr())
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept connection:", err.Error())
			continue
		}
		log.Printf("Connection: client=%v\n", conn.RemoteAddr())
		go handleConn(serverConfig, conn)
	}
}

func handleConn(serverConfig *ssh.ServerConfig, conn net.Conn) {
	defer conn.Close()
	_, channels, requests, err := ssh.NewServerConn(conn, serverConfig)
	if err != nil {
		log.Println("Failed to establish SSH connection:", err.Error())
		return
	}
	log.Printf("Established SSH connection: client=%v\n", conn.RemoteAddr())
	go request.Handle(conn.RemoteAddr(), "global", requests)
	for newChannel := range channels {
		go channel.Handle(conn.RemoteAddr(), newChannel)
	}
	log.Printf("Disconnect: client=%v\n", conn.RemoteAddr())
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"gopkg.in/urfave/cli.v1"
	"log"
	"os"
	"regexp"
	"strings"
)

type ArgConfig struct {
	AwsRegion       string
	SecurityGroupId string
	// Port or ICMP Type
	From int
	// Port or ICMP Type
	To                   int
	Protocol             string
	ThisEc2Instance      bool
	AutoScalingGroupName string
}

func parseFlags(c *cli.Context) *ArgConfig {
	var result ArgConfig
	result.AwsRegion = c.String("r")
	re := regexp.MustCompile("^sg-([0-9a-z]{8})$")
	if re.MatchString(c.String("s")) == false {
		log.Printf("Error: -s must be specified as a Security Group ID. Example: sg-asdf1234\n\n")
		cli.ShowAppHelp(c)
		os.Exit(1)
	} else {
		result.SecurityGroupId = c.String("s")
	}
	if c.String("p") != "tcp" && c.String("p") != "udp" && c.String("p") != "icmp" {
		log.Printf("Error: -p must be one of 'tcp', 'udp', or 'icmp'.\n\n")
		cli.ShowAppHelp(c)
		os.Exit(1)
	} else {
		result.Protocol = c.String("p")
	}
	if result.Protocol == "tcp" || result.Protocol == "udp" {
		if c.Int("f") < 1 || c.Int("f") > 65535 {
			log.Printf("Error: -f must be a port number between 1 and 65535.\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.From = c.Int("f")
		}
		if c.Int("t") < 1 || c.Int("t") > 65535 {
			log.Printf("Error: -t must be a port number between 1 and 65535.\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.To = c.Int("t")
		}
	} else {
		if c.Int("f") < 0 || c.Int("f") > 255 {
			log.Printf("Error: -f must be an ICMP type between 0 and 255. See http://www.iana.org/assignments/icmp-parameters/icmp-parameters.xhtml\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.From = c.Int("f")
		}
		if c.Int("t") < 0 || c.Int("t") > 255 {
			log.Printf("Error: -t must be an ICMP type between 0 and 255. See http://www.iana.org/assignments/icmp-parameters/icmp-parameters.xhtml\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.To = c.Int("t")
		}
	}
	result.ThisEc2Instance = c.Bool("i")
	result.AutoScalingGroupName = c.String("a")
	return &result
}

func main() {
	app := cli.NewApp()
	app.Name = "ec2-sg-mangler"
	app.Version = "0.1.0"
	app.Usage = "Helper utility to manage the EC2 instance public IPs in an AWS Security Group"
	app.Action = func(c *cli.Context) error {
		arg_config := parseFlags(c)

		ec2metadata_client := ec2metadata.New(session.New())
		err := setAwsRegion(ec2metadata_client, arg_config)
		if err != nil {
			log.Fatalf(err.Error())
		}

		// Reusable config session object for AWS services with current region attached.
		aws_config_session := session.New(&aws.Config{Region: aws.String(arg_config.AwsRegion)})

		asg_client := autoscaling.New(aws_config_session)
		ec2_client := ec2.New(aws_config_session)

		var proposed_ips []string
		if len(arg_config.AutoScalingGroupName) > 0 {
			log.Printf("Fetching list of Auto Scaling Group IPs...\n")
			proposed_ips, err = getAsgInstancePublicIps(asg_client, ec2_client, arg_config.AutoScalingGroupName)
			if err != nil {
				log.Fatalf(err.Error())
			}
		}
		var this_instance_ip string
		if arg_config.ThisEc2Instance == true {
			log.Printf("Fetching the public IP of this EC2 instance...\n")
			this_instance_ip, err = getThisInstancePublicIp(ec2metadata_client)
			if err != nil {
				log.Fatalf(err.Error())
			}
			proposed_ips = append(proposed_ips, this_instance_ip)
		}
		proposed_ips = removeSliceDuplicates(proposed_ips)

		log.Printf("Fetching the IPs currently attached to Security Group...\n")
		sg_ips, err := getCurrentMatchingSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol)
		if err != nil {
			log.Fatalf(err.Error())
		}

		sg_actions := reconcileIps(sg_ips, proposed_ips)
		log.Printf("Commencing SG changes (if any)...\n")
		if len(sg_actions.Add) > 0 {
			err = doAddSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol, sg_actions.Add)
			if err != nil {
				log.Fatalf(err.Error())
			}
			log.Printf("Added Ingress IPs %s to SG %s.\n", strings.Join(sg_actions.Add, ", "), arg_config.SecurityGroupId)
		}
		if len(sg_actions.Remove) > 0 {
			err = doRemoveSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol, sg_actions.Remove)
			if err != nil {
				log.Fatalf(err.Error())
			}
			log.Printf("Removed Ingress IPs %s to SG %s.\n", strings.Join(sg_actions.Remove, ", "), arg_config.SecurityGroupId)
		}
		log.Printf("All done.")
		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "r",
			Value: "",
			Usage: "AWS Region",
		},
		cli.StringFlag{
			Name:  "s",
			Value: "",
			Usage: "Security Group ID (must exist already)",
		},
		cli.IntFlag{
			Name:  "f",
			Value: 1,
			Usage: "From Port Number OR ICMP Type (-1 for all ICMP types)",
		},
		cli.IntFlag{
			Name:  "t",
			Value: 1,
			Usage: "To Port Number OR ICMP Type (-1 for all ICMP types )",
		},
		cli.StringFlag{
			Name:  "p",
			Value: "tcp",
			Usage: "IP Protocol Name (may be one of 'tcp', 'udp', or 'icmp' currently)",
		},
		cli.BoolFlag{
			Name:  "i",
			Usage: "Add this EC2 instance public IP? Not required if this instance is part of the nominated ASG",
		},
		cli.StringFlag{
			Name:  "a",
			Value: "",
			Usage: "Add the public IPs for all EC2 instances in this ASG Name",
		},
	}

	app.Run(os.Args)
}
<commit_msg>bump version<commit_after>package main

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"gopkg.in/urfave/cli.v1"
	"log"
	"os"
	"regexp"
	"strings"
)

type ArgConfig struct {
	AwsRegion       string
	SecurityGroupId string
	// Port or ICMP Type
	From int
	// Port or ICMP Type
	To                   int
	Protocol             string
	ThisEc2Instance      bool
	AutoScalingGroupName string
}

func parseFlags(c *cli.Context) *ArgConfig {
	var result ArgConfig
	result.AwsRegion = c.String("r")
	re := regexp.MustCompile("^sg-([0-9a-z]{8})$")
	if re.MatchString(c.String("s")) == false {
		log.Printf("Error: -s must be specified as a Security Group ID. Example: sg-asdf1234\n\n")
		cli.ShowAppHelp(c)
		os.Exit(1)
	} else {
		result.SecurityGroupId = c.String("s")
	}
	if c.String("p") != "tcp" && c.String("p") != "udp" && c.String("p") != "icmp" {
		log.Printf("Error: -p must be one of 'tcp', 'udp', or 'icmp'.\n\n")
		cli.ShowAppHelp(c)
		os.Exit(1)
	} else {
		result.Protocol = c.String("p")
	}
	if result.Protocol == "tcp" || result.Protocol == "udp" {
		if c.Int("f") < 1 || c.Int("f") > 65535 {
			log.Printf("Error: -f must be a port number between 1 and 65535.\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.From = c.Int("f")
		}
		if c.Int("t") < 1 || c.Int("t") > 65535 {
			log.Printf("Error: -t must be a port number between 1 and 65535.\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.To = c.Int("t")
		}
	} else {
		if c.Int("f") < 0 || c.Int("f") > 255 {
			log.Printf("Error: -f must be an ICMP type between 0 and 255. See http://www.iana.org/assignments/icmp-parameters/icmp-parameters.xhtml\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.From = c.Int("f")
		}
		if c.Int("t") < 0 || c.Int("t") > 255 {
			log.Printf("Error: -t must be an ICMP type between 0 and 255. See http://www.iana.org/assignments/icmp-parameters/icmp-parameters.xhtml\n\n")
			cli.ShowAppHelp(c)
			os.Exit(1)
		} else {
			result.To = c.Int("t")
		}
	}
	result.ThisEc2Instance = c.Bool("i")
	result.AutoScalingGroupName = c.String("a")
	return &result
}

func main() {
	app := cli.NewApp()
	app.Name = "ec2-sg-mangler"
	app.Version = "0.2.0"
	app.Usage = "Helper utility to manage the EC2 instance public IPs in an AWS Security Group"
	app.Action = func(c *cli.Context) error {
		arg_config := parseFlags(c)

		ec2metadata_client := ec2metadata.New(session.New())
		err := setAwsRegion(ec2metadata_client, arg_config)
		if err != nil {
			log.Fatalf(err.Error())
		}

		// Reusable config session object for AWS services with current region attached.
		aws_config_session := session.New(&aws.Config{Region: aws.String(arg_config.AwsRegion)})

		asg_client := autoscaling.New(aws_config_session)
		ec2_client := ec2.New(aws_config_session)

		var proposed_ips []string
		if len(arg_config.AutoScalingGroupName) > 0 {
			log.Printf("Fetching list of Auto Scaling Group IPs...\n")
			proposed_ips, err = getAsgInstancePublicIps(asg_client, ec2_client, arg_config.AutoScalingGroupName)
			if err != nil {
				log.Fatalf(err.Error())
			}
		}
		var this_instance_ip string
		if arg_config.ThisEc2Instance == true {
			log.Printf("Fetching the public IP of this EC2 instance...\n")
			this_instance_ip, err = getThisInstancePublicIp(ec2metadata_client)
			if err != nil {
				log.Fatalf(err.Error())
			}
			proposed_ips = append(proposed_ips, this_instance_ip)
		}
		proposed_ips = removeSliceDuplicates(proposed_ips)

		log.Printf("Fetching the IPs currently attached to Security Group...\n")
		sg_ips, err := getCurrentMatchingSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol)
		if err != nil {
			log.Fatalf(err.Error())
		}

		sg_actions := reconcileIps(sg_ips, proposed_ips)
		log.Printf("Commencing SG changes (if any)...\n")
		if len(sg_actions.Add) > 0 {
			err = doAddSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol, sg_actions.Add)
			if err != nil {
				log.Fatalf(err.Error())
			}
			log.Printf("Added Ingress IPs %s to SG %s.\n", strings.Join(sg_actions.Add, ", "), arg_config.SecurityGroupId)
		}
		if len(sg_actions.Remove) > 0 {
			err = doRemoveSgIps(ec2_client, arg_config.SecurityGroupId, arg_config.From, arg_config.To, arg_config.Protocol, sg_actions.Remove)
			if err != nil {
				log.Fatalf(err.Error())
			}
			log.Printf("Removed Ingress IPs %s to SG %s.\n", strings.Join(sg_actions.Remove, ", "), arg_config.SecurityGroupId)
		}
		log.Printf("All done.")
		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "r",
			Value: "",
			Usage: "AWS Region",
		},
		cli.StringFlag{
			Name:  "s",
			Value: "",
			Usage: "Security Group ID (must exist already)",
		},
		cli.IntFlag{
			Name:  "f",
			Value: 1,
			Usage: "From Port Number OR ICMP Type (-1 for all ICMP types)",
		},
		cli.IntFlag{
			Name:  "t",
			Value: 1,
			Usage: "To Port Number OR ICMP Type (-1 for all ICMP types )",
		},
		cli.StringFlag{
			Name:  "p",
			Value: "tcp",
			Usage: "IP Protocol Name (may be one of 'tcp', 'udp', or 'icmp' currently)",
		},
		cli.BoolFlag{
			Name:  "i",
			Usage: "Add this EC2 instance public IP? Not required if this instance is part of the nominated ASG",
		},
		cli.StringFlag{
			Name:  "a",
			Value: "",
			Usage: "Add the public IPs for all EC2 instances in this ASG Name",
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"time"

	"github.com/Liamraystanley/marill/domfinder"
	"github.com/Liamraystanley/marill/scraper"
	"github.com/urfave/cli"
)

// these /SHOULD/ be defined during the make process. not always however.
var version, commithash, compiledate = "", "", ""

type outputConfig struct {
	noColors   bool
	printDebug bool
	ignoreStd  bool
	logFile    string
}

type scanConfig struct {
	cores int
}

type appConfig struct {
	printUrls bool
}

type config struct {
	app  appConfig
	scan scanConfig
	out  outputConfig
}

var conf config
var out = Output{}

func statsLoop(done <-chan struct{}) {
	mem := &runtime.MemStats{}
	var numRoutines, numCPU int
	var load5, load10, load15 float32

	for {
		select {
		case <-done:
			return
		default:
			runtime.ReadMemStats(mem)
			numRoutines = runtime.NumGoroutine()
			numCPU = runtime.NumCPU()

			if contents, err := ioutil.ReadFile("/proc/loadavg"); err == nil {
				fmt.Sscanf(string(contents), "%f %f %f %*s %*d", &load5, &load10, &load15)
			}

			logger.Printf(
				"allocated mem: %dM, sys: %dM, threads: %d, cores: %d load5: %.2f load10: %.2f load15: %.2f",
				mem.Alloc/1024/1024, mem.Sys/1024/1024, numRoutines, numCPU, load5, load10, load15)

			time.Sleep(2 * time.Second)
		}
	}
}

func numCores() {
	if conf.scan.cores == 0 {
		if runtime.NumCPU() == 1 {
			conf.scan.cores = 1
		} else {
			conf.scan.cores = runtime.NumCPU() / 2
		}
	} else if conf.scan.cores > runtime.NumCPU() {
		logger.Printf("warning: using %d cores, which is more than the amount of cores", conf.scan.cores)
		out.Printf("{yellow}warning: using %d cores, which is more than the amount of cores on the server!{c}\n", conf.scan.cores)

		// set it to the amount of cores on the server. go will do this regardless, so.
		conf.scan.cores = runtime.NumCPU()
		logger.Printf("limiting number of cores to %d", conf.scan.cores)
		out.Printf("limiting number of cores to %d\n", conf.scan.cores)
	}

	runtime.GOMAXPROCS(conf.scan.cores)
	logger.Printf("using %d cores (max %d)", conf.scan.cores, runtime.NumCPU())

	return
}

func printUrls() error {
	finder := &domfinder.Finder{Log: logger}
	if err := finder.GetWebservers(); err != nil {
		return fmt.Errorf("unable to get process list: %s", err)
	}

	if err := finder.GetDomains(); err != nil {
		return fmt.Errorf("unable to auto-fetch domain list: %s", err)
	}

	for _, domain := range finder.Domains {
		out.Printf("{blue}%-40s{c} {green}%s{c}\n", domain.URL, domain.IP)
	}

	return nil
}

func run() {
	logger.Println("checking for running webservers...")

	finder := &domfinder.Finder{Log: logger}
	if err := finder.GetWebservers(); err != nil {
		logger.Fatalf("unable to get process list: %s", err)
	}

	if out := ""; len(finder.Procs) > 0 {
		for _, proc := range finder.Procs {
			out += fmt.Sprintf("[%s:%s] ", proc.Name, proc.PID)
		}
		logger.Printf("found %d procs matching a webserver: %s", len(finder.Procs), out)
	}

	// start crawling for domains
	if err := finder.GetDomains(); err != nil {
		logger.Fatalf("unable to auto-fetch domain list: %s", err)
	}

	logger.Printf("found %d domains on webserver %s (exe: %s, pid: %s)", len(finder.Domains), finder.MainProc.Name, finder.MainProc.Exe, finder.MainProc.PID)

	tmplist := []*scraper.Domain{}
	for _, domain := range finder.Domains {
		tmplist = append(tmplist, &scraper.Domain{URL: domain.URL, IP: domain.IP})
	}
	crawler := &scraper.Crawler{Log: logger, Domains: tmplist}
	crawler.Crawl()
}

func main() {
	defer closeLogger() // ensure we're cleaning up the logger if there is one

	cli.VersionPrinter = func(c *cli.Context) {
		if version != "" && commithash != "" && compiledate != "" {
			fmt.Printf("version %s, revision %s (%s)\n", version, commithash, compiledate)
		} else if commithash != "" && compiledate != "" {
			fmt.Printf("revision %s (%s)\n", commithash, compiledate)
		} else if version != "" {
			fmt.Printf("version %s\n", version)
		} else {
			fmt.Println("version unknown")
		}
	}

	app := cli.NewApp()

	app.Name = "marill"

	if version != "" && commithash != "" {
		app.Version = fmt.Sprintf("%s, git revision %s", version, commithash)
	} else if version != "" {
		app.Version = version
	} else if commithash != "" {
		app.Version = "git revision " + commithash
	}

	app.Authors = []cli.Author{
		cli.Author{
			Name:  "Liam Stanley",
			Email: "me@liamstanley.io",
		},
	}
	app.Compiled = time.Now()
	app.Usage = "Automated website testing utility"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:        "printurls",
			Usage:       "Print the list of urls as if they were going to be scanned",
			Destination: &conf.app.printUrls,
		},
		cli.BoolFlag{
			Name:        "debug, d",
			Usage:       "Print debugging information to stdout",
			Destination: &conf.out.printDebug,
		},
		cli.BoolFlag{
			Name:        "quiet, q",
			Usage:       "Dont't print regular stdout messages",
			Destination: &conf.out.ignoreStd,
		},
		cli.StringFlag{
			Name:        "log-file",
			Usage:       "Log debugging information to `logfile`",
			Destination: &conf.out.logFile,
		},
		cli.IntFlag{
			Name:        "cores",
			Usage:       "Use `n` cores to fetch data (0 being server cores/2)",
			Destination: &conf.scan.cores,
		},
	}

	app.Action = func(c *cli.Context) error {
		// initialize the logger. ensure this only occurs after the cli args are
		// pulled.
		initLogger()

		// initialize some form of max go procs
		numCores()

		// initialize the stats data
		done := make(chan struct{}, 1)
		go statsLoop(done)

		// close the stats data goroutine when we're complete.
		defer func() {
			done <- struct{}{}
		}()

		if conf.app.printUrls {
			if err := printUrls(); err != nil {
				fmt.Printf("err: %s", err)
				return err
			}

			os.Exit(0)
		}

		run()

		return nil
	}

	app.Run(os.Args)
}
<commit_msg>allow filtering of http vs https urls<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"time"

	"github.com/Liamraystanley/marill/domfinder"
	"github.com/Liamraystanley/marill/scraper"
	"github.com/urfave/cli"
)

// these /SHOULD/ be defined during the make process. not always however.
var version, commithash, compiledate = "", "", ""

type outputConfig struct {
	noColors   bool
	printDebug bool
	ignoreStd  bool
	logFile    string
}

type scanConfig struct {
	cores       int
	ignorehttp  bool
	ignorehttps bool
}

type appConfig struct {
	printUrls bool
}

type config struct {
	app  appConfig
	scan scanConfig
	out  outputConfig
}

var conf config
var out = Output{}

func statsLoop(done <-chan struct{}) {
	mem := &runtime.MemStats{}
	var numRoutines, numCPU int
	var load5, load10, load15 float32

	for {
		select {
		case <-done:
			return
		default:
			runtime.ReadMemStats(mem)
			numRoutines = runtime.NumGoroutine()
			numCPU = runtime.NumCPU()

			if contents, err := ioutil.ReadFile("/proc/loadavg"); err == nil {
				fmt.Sscanf(string(contents), "%f %f %f %*s %*d", &load5, &load10, &load15)
			}

			logger.Printf(
				"allocated mem: %dM, sys: %dM, threads: %d, cores: %d load5: %.2f load10: %.2f load15: %.2f",
				mem.Alloc/1024/1024, mem.Sys/1024/1024, numRoutines, numCPU, load5, load10, load15)

			time.Sleep(2 * time.Second)
		}
	}
}

func numCores() {
	if conf.scan.cores == 0 {
		if runtime.NumCPU() == 1 {
			conf.scan.cores = 1
		} else {
			conf.scan.cores = runtime.NumCPU() / 2
		}
	} else if conf.scan.cores > runtime.NumCPU() {
		logger.Printf("warning: using %d cores, which is more than the amount of cores", conf.scan.cores)
		out.Printf("{yellow}warning: using %d cores, which is more than the amount of cores on the server!{c}\n", conf.scan.cores)

		// set it to the amount of cores on the server. go will do this regardless, so.
		conf.scan.cores = runtime.NumCPU()
		logger.Printf("limiting number of cores to %d", conf.scan.cores)
		out.Printf("limiting number of cores to %d\n", conf.scan.cores)
	}

	runtime.GOMAXPROCS(conf.scan.cores)
	logger.Printf("using %d cores (max %d)", conf.scan.cores, runtime.NumCPU())

	return
}

func printUrls() error {
	finder := &domfinder.Finder{Log: logger}
	if err := finder.GetWebservers(); err != nil {
		return fmt.Errorf("unable to get process list: %s", err)
	}

	if err := finder.GetDomains(); err != nil {
		return fmt.Errorf("unable to auto-fetch domain list: %s", err)
	}

	finder.Filter(domfinder.DomainFilter{
		IgnoreHTTP:  conf.scan.ignorehttp,
		IgnoreHTTPS: conf.scan.ignorehttps,
	})

	for _, domain := range finder.Domains {
		out.Printf("{blue}%-40s{c} {green}%s{c}\n", domain.URL, domain.IP)
	}

	return nil
}

func run() {
	logger.Println("checking for running webservers...")

	finder := &domfinder.Finder{Log: logger}
	if err := finder.GetWebservers(); err != nil {
		logger.Fatalf("unable to get process list: %s", err)
	}

	if out := ""; len(finder.Procs) > 0 {
		for _, proc := range finder.Procs {
			out += fmt.Sprintf("[%s:%s] ", proc.Name, proc.PID)
		}
		logger.Printf("found %d procs matching a webserver: %s", len(finder.Procs), out)
	}

	// start crawling for domains
	if err := finder.GetDomains(); err != nil {
		logger.Fatalf("unable to auto-fetch domain list: %s", err)
	}

	finder.Filter(domfinder.DomainFilter{
		IgnoreHTTP:  conf.scan.ignorehttp,
		IgnoreHTTPS: conf.scan.ignorehttps,
	})

	logger.Printf("found %d domains on webserver %s (exe: %s, pid: %s)", len(finder.Domains), finder.MainProc.Name, finder.MainProc.Exe, finder.MainProc.PID)

	tmplist := []*scraper.Domain{}
	for _, domain := range finder.Domains {
		tmplist = append(tmplist, &scraper.Domain{URL: domain.URL, IP: domain.IP})
	}
	crawler := &scraper.Crawler{Log: logger, Domains: tmplist}
	crawler.Crawl()
}

func main() {
	defer closeLogger() // ensure we're cleaning up the logger if there is one

	cli.VersionPrinter = func(c *cli.Context) {
		if version != "" && commithash != "" && compiledate != "" {
			fmt.Printf("version %s, revision %s (%s)\n", version, commithash, compiledate)
		} else if commithash != "" && compiledate != "" {
			fmt.Printf("revision %s (%s)\n", commithash, compiledate)
		} else if version != "" {
			fmt.Printf("version %s\n", version)
		} else {
			fmt.Println("version unknown")
		}
	}

	app := cli.NewApp()

	app.Name = "marill"

	if version != "" && commithash != "" {
		app.Version = fmt.Sprintf("%s, git revision %s", version, commithash)
	} else if version != "" {
		app.Version = version
	} else if commithash != "" {
		app.Version = "git revision " + commithash
	}

	app.Authors = []cli.Author{
		cli.Author{
			Name:  "Liam Stanley",
			Email: "me@liamstanley.io",
		},
	}
	app.Compiled = time.Now()
	app.Usage = "Automated website testing utility"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:        "printurls",
			Usage:       "Print the list of urls as if they were going to be scanned",
			Destination: &conf.app.printUrls,
		},
		cli.BoolFlag{
			Name:        "debug, d",
			Usage:       "Print debugging information to stdout",
			Destination: &conf.out.printDebug,
		},
		cli.BoolFlag{
			Name:        "quiet, q",
			Usage:       "Dont't print regular stdout messages",
			Destination: &conf.out.ignoreStd,
		},
		cli.StringFlag{
			Name:        "log-file",
			Usage:       "Log debugging information to `logfile`",
			Destination: &conf.out.logFile,
		},
		cli.IntFlag{
			Name:        "cores",
			Usage:       "Use `n` cores to fetch data (0 being server cores/2)",
			Destination: &conf.scan.cores,
		},
		cli.BoolFlag{
			Name:        "ignore-http",
			Usage:       "Ignore http-based URLs in domain finder",
			Destination: &conf.scan.ignorehttp,
		},
		cli.BoolFlag{
			Name:        "ignore-https",
			Usage:       "Ignore https-based URLs in domain finder",
			Destination: &conf.scan.ignorehttps,
		},
	}

	app.Action = func(c *cli.Context) error {
		// initialize the logger. ensure this only occurs after the cli args are
		// pulled.
		initLogger()

		// initialize some form of max go procs
		numCores()

		// initialize the stats data
		done := make(chan struct{}, 1)
		go statsLoop(done)

		// close the stats data goroutine when we're complete.
		defer func() {
			done <- struct{}{}
		}()

		if conf.app.printUrls {
			if err := printUrls(); err != nil {
				fmt.Printf("err: %s", err)
				return err
			}

			os.Exit(0)
		}

		run()

		return nil
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"hash/crc32"
	"io"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
)

const (
	membersCount = 10
	emailDomain  = "example.com"
)

var (
	nameFirstLetters = []byte("abcdefghijklmnopqrstuvwz")
	lastNames        = []string{
		"frerichs", "hengst", "kloth", "orth", "popp", "steigerwald", "kley", "voegele", "ostermeier",
		"schley", "eckel", "kaeser", "holzner", "klostermann", "ostertag", "weishaupt", "gotte", "feiler", "mast",
		"hielscher", "haertel", "moeller", "kramer", "stief", "kissel", "gottschlich", "drechsler", "lucas", "till",
		"koehne", "pfitzner", "sydow", "liedtke", "franz", "flick", "menz", "hertel", "kreuzer", "kleinert", "zeidler",
	}
)

type myHandler struct{}

type InfrastructureAccount struct {
	Id          string `json:"id"`
	Name        string `json:"name"`
	Provider    string `json:"provider"`
	Type        string `json:"type"`
	Description string `json:"description"`
	Owner       string `json:"owner"`
	OwnerDn     string `json:"owner_dn"`
	Disabled    bool   `json:"disabled"`
}

type Team struct {
	Dn           string   `json:"dn"`
	Id           string   `json:"id"`
	TeamName     string   `json:"id_name"`
	TeamId       string   `json:"team_id"`
	Type         string   `json:"type"`
	FullName     string   `json:"name"`
	Aliases      []string `json:"alias,omitempty"`
	Mails        []string `json:"mail"`
	Members      []string `json:"member"`
	CostCenter   string   `json:"cost_center,omitempty"`
	DeliveryLead string   `json:"delivery_lead,omitempty"`
	ParentTeamId string   `json:"parent_team_id,omitempty"`

	InfrastructureAccounts []InfrastructureAccount `json:"infrastructure-accounts"`
}

func main() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM) // Push signals into channel

	if len(os.Args) < 2 {
		fmt.Printf("Usage %s {port}\n", os.Args[0])
		os.Exit(1)
	}

	port, err := strconv.Atoi(os.Args[1])
	if err != nil {
		fmt.Printf("Port must be numeric\n")
		os.Exit(1)
	}

	server := http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: &myHandler{},
	}

	go func() {
		err = server.ListenAndServe()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}()
	fmt.Printf("Listening on %d\n", port)

	sig := <-sigs
	fmt.Printf("Got %+v signal. Shutting down\n", sig)
}

func members() []string {
	res := make([]string, membersCount)
	for i := 0; i < membersCount; i++ {
		name := nameFirstLetters[rand.Intn(len(nameFirstLetters))]
		res[i] = string(name) + lastNames[rand.Intn(len(lastNames))]
	}

	return res
}

func testTeam(teamName string, w http.ResponseWriter) {
	id := crc32.ChecksumIEEE([]byte(teamName))
	rand.Seed(int64(id))

	teamType := "official"
	if id%2 == 0 {
		teamType = "virtual"
	}

	team := Team{
		Id:       teamName,
		TeamId:   strconv.Itoa(int(id)),
		TeamName: strings.ToUpper(teamName),
		Members:  members(),
		FullName: strings.Title(teamName),
		Mails:    []string{fmt.Sprintf("%s@%s", teamName, emailDomain)},
		Type:     teamType,
	}

	m, err := json.Marshal(team)
	if err != nil {
		fmt.Fprintf(w, "Can't marshal: %s", err)
	}
	w.Header().Set("Content-Type", "application/json")

	w.Write(m)
}

func badRequest(w http.ResponseWriter) {
	w.WriteHeader(http.StatusBadRequest)
	w.Write([]byte("Bad request"))
}

func (*myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Printf("Request: %s\n", r.URL.Path)
	if !strings.HasPrefix(r.URL.Path, "/teams") {
		badRequest(w)
		return
	}
	var teamName string

	parts := strings.Split(r.URL.Path, "/")
	if len(parts) > 1 {
		teamName = parts[2]
	}

	if teamName == "" {
		badRequest(w)
		return
	}

	testTeam(teamName, w)
}
<commit_msg>make use of log<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"hash/crc32"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
)

const (
	membersCount = 10
	emailDomain  = "example.com"
)

var (
	nameFirstLetters = []byte("abcdefghijklmnopqrstuvwz")
	lastNames        = []string{
		"frerichs", "hengst", "kloth", "orth", "popp", "steigerwald", "kley", "voegele", "ostermeier",
		"schley", "eckel", "kaeser", "holzner", "klostermann", "ostertag", "weishaupt", "gotte", "feiler", "mast",
		"hielscher", "haertel", "moeller", "kramer", "stief", "kissel", "gottschlich", "drechsler", "lucas", "till",
		"koehne", "pfitzner", "sydow", "liedtke", "franz", "flick", "menz", "hertel", "kreuzer", "kleinert", "zeidler",
	}
)

type myHandler struct{}

type InfrastructureAccount struct {
	Id          string `json:"id"`
	Name        string `json:"name"`
	Provider    string `json:"provider"`
	Type        string `json:"type"`
	Description string `json:"description"`
	Owner       string `json:"owner"`
	OwnerDn     string `json:"owner_dn"`
	Disabled    bool   `json:"disabled"`
}

type Team struct {
	Dn           string   `json:"dn"`
	Id           string   `json:"id"`
	TeamName     string   `json:"id_name"`
	TeamId       string   `json:"team_id"`
	Type         string   `json:"type"`
	FullName     string   `json:"name"`
	Aliases      []string `json:"alias,omitempty"`
	Mails        []string `json:"mail"`
	Members      []string `json:"member"`
	CostCenter   string   `json:"cost_center,omitempty"`
	DeliveryLead string   `json:"delivery_lead,omitempty"`
	ParentTeamId string   `json:"parent_team_id,omitempty"`

	InfrastructureAccounts []InfrastructureAccount `json:"infrastructure-accounts"`
}

func main() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM) // Push signals into channel

	if len(os.Args) < 2 {
		fmt.Printf("Usage %s {port}\n", os.Args[0])
		os.Exit(1)
	}

	port, err := strconv.Atoi(os.Args[1])
	if err != nil {
		fmt.Printf("Port must be numeric\n")
		os.Exit(1)
	}

	server := http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: &myHandler{},
	}

	go func() {
		err := server.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()
	log.Printf("Listening on %d\n", port)

	sig := <-sigs
	log.Printf("Got %+v signal. Shutting down\n", sig)
}

func members() []string {
	res := make([]string, membersCount)
	for i := 0; i < membersCount; i++ {
		name := nameFirstLetters[rand.Intn(len(nameFirstLetters))]
		res[i] = string(name) + lastNames[rand.Intn(len(lastNames))]
	}

	return res
}

func testTeam(teamName string, w http.ResponseWriter) {
	id := crc32.ChecksumIEEE([]byte(teamName))
	rand.Seed(int64(id))

	teamType := "official"
	if id%2 == 0 {
		teamType = "virtual"
	}

	team := Team{
		Id:       teamName,
		TeamId:   strconv.Itoa(int(id)),
		TeamName: strings.ToUpper(teamName),
		Members:  members(),
		FullName: strings.Title(teamName),
		Mails:    []string{fmt.Sprintf("%s@%s", teamName, emailDomain)},
		Type:     teamType,
	}

	m, err := json.Marshal(team)
	if err != nil {
		fmt.Fprintf(w, "Can't marshal: %s", err)
	}
	w.Header().Set("Content-Type", "application/json")

	w.Write(m)
}

func badRequest(w http.ResponseWriter) {
	w.WriteHeader(http.StatusBadRequest)
	w.Write([]byte("Bad request"))
}

func (*myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("Request: %s\n", r.URL.Path)
	if !strings.HasPrefix(r.URL.Path, "/teams") {
		badRequest(w)
		return
	}
	var teamName string

	parts := strings.Split(r.URL.Path, "/")
	if len(parts) > 1 {
		teamName = parts[2]
	}

	if teamName == "" {
		badRequest(w)
		return
	}

	testTeam(teamName, w)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	alexandria "github.com/maiconio/alexandria/orm"
)

//Book is on the table
type Book struct {
	ID    int
	Name  string
	Pages int
}

func main() {
	//connect to Postgres
	orm, scream := alexandria.ConnectToPostgres()
	if scream != nil {
		panic(scream)
	}

	//create the orm handler for Book objects
	ormBooks := orm.Handle(Book{})

	//Create Table
	ormBooks.CreateTable()

	//Insert/update
	ormBooks.Save(Book{Name: "The book is on the table", Pages: 198})

	//Select
	ormBooks.Find().Where("pages > 0")
	ormBooks.Find().ByID(9)
	ormBooks.Find().All()

	//Delete
	ormBooks.Delete().Where("id=9")
	ormBooks.Delete().ByID(9)
	ormBooks.Delete().All()

	//Drop Table
	ormBooks.DropTable()
}
<commit_msg>Update API design.<commit_after>package main

import (
	"fmt"
	alexandria "github.com/maiconio/alexandria/orm"
)

//Book is on the table
type Book struct {
	ID    int
	Name  string
	Pages int
}

func main() {
	//connect to Postgres
	orm, scream := alexandria.ConnectToPostgres()
	if scream != nil {
		panic(scream)
	}

	//create the orm handler for Book objects
	ormBooks := orm.NewHandler(Book{})

	//Create Table
	ormBooks.DropTable()
	ormBooks.CreateTable()
	err := ormBooks.Insert(Book{Name: "The book is on the table", Pages: 198})
	fmt.Printf("err: %v\n", err)

	/*
		//Insert/update
		ormBooks.Insert(Book{Name: "The book is on the table", Pages: 198})

		//Update
		ormBooks.Update().Where("pages > 0")
		ormBooks.Update().ByID(9)
		ormBooks.Update().All()

		//Select
		ormBooks.Select().Where("pages > 0")
		ormBooks.Select().ByID(9)
		ormBooks.Select().All()

		//Delete
		ormBooks.Delete().Where("id=9")
		ormBooks.Delete().ByID(9)
		ormBooks.Delete().All()

		//Drop Table
		ormBooks.DropTable()
	*/
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"takiones.com/goprompt/colors"
)

func getPythonVirtualEnv() string {
	virtualEnv, ve := os.LookupEnv("VIRTUAL_ENV")
	if ve {
		ave := strings.Split(virtualEnv, "/")
		virtualEnv = fmt.Sprintf("(%s) ", ave[len(ave)-1])
	}
	return virtualEnv
}

func main() {
	//Get basicinfo
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Unable to get current path", err)
	}
	user := os.Getenv("USER")
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Unable to get hostname", err)
	}
	//Get Python VirtualEnv info
	virtualEnv := getPythonVirtualEnv()

	//Format output
	userInfo := colors.Format(user+"@"+hostname, "bold", "green")
	pwdInfo := colors.Format(pwd, "bold", "blue")
	virtualEnvInfo := colors.Format(virtualEnv, "blue")

	fmt.Printf("%s[%s  %s]$ ", virtualEnvInfo, userInfo, pwdInfo)
}
<commit_msg>initial git<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	git2go "gopkg.in/libgit2/git2go.v26"

	"takiones.com/goprompt/colors"
)

func getPythonVirtualEnv() string {
	virtualEnv, ve := os.LookupEnv("VIRTUAL_ENV")
	if ve {
		ave := strings.Split(virtualEnv, "/")
		virtualEnv = fmt.Sprintf("(%s) ", ave[len(ave)-1])
	}
	return virtualEnv
}

func main() {
	//Get basicinfo
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Unable to get current path", err)
	}
	user := os.Getenv("USER")
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Unable to get hostname", err)
	}
	//Get Python VirtualEnv info
	virtualEnv := getPythonVirtualEnv()

	//Get git information
	_ = git2go.Repository{}

	discover:= git2go.Discover(".",false, ["/","."] )

	//Format output
	userInfo := colors.Format(user+"@"+hostname, "bold", "green")
	pwdInfo := colors.Format(pwd, "bold", "blue")
	virtualEnvInfo := colors.Format(virtualEnv, "blue")

	fmt.Printf("%s[%s %s]$ ", virtualEnvInfo, userInfo, pwdInfo)
}
<|endoftext|>


<commit_before>package main

import "github.com/howeyc/fsnotify"
import "log"
import "flag"
import "fmt"
import "os"
import "os/exec"
import "regexp"

var versionStr = "0.1.0"

var goCommandSet = map[string][]string{
	"test":    []string{"go", "test"},
	"install": []string{"go", "install"},
	"build":   []string{"go", "build"},
}

func main() {
	var helpFlag = flag.Bool("h", false, "help")
	var buildFlag = flag.Bool("b", true, "build, default behavior")
	var testFlag = flag.Bool("t", false, "test")
	var installFlag = flag.Bool("i", false, "install")
	var versionFlag = flag.Bool("v", false, "version")
	var xFlag = flag.Bool("x", false, "show verbose command")

	flag.Parse()
	args := flag.Args()

	if *helpFlag {
		fmt.Println("Usage: gomon [options] [dir] [-- command]")
		fmt.Println("   -b build")
		fmt.Println("   -t test")
		fmt.Println("   -i install")
		fmt.Println("   -x show verbose command")
		fmt.Println("   -h help")
		os.Exit(0)
	}
	if *versionFlag {
		fmt.Printf("gomon %s\n", versionStr)
		os.Exit(0)
	}

	var dirs = []string{}
	var cmds = []string{}

	var takeDir = true
	for _, a := range args {
		var exists, _ = FileExists(a)
		if a == "--" || !exists {
			takeDir = false
			continue
		}
		if takeDir {
			dirs = append(dirs, a)
		} else {
			cmds = append(cmds, a)
		}
	}

	if len(cmds) == 0 {
		if *testFlag {
			cmds = goCommandSet["test"]
		} else if *buildFlag {
			cmds = goCommandSet["build"]
		} else if *installFlag {
			cmds = goCommandSet["install"]
		}
		if *xFlag && len(cmds) > 0 {
			cmds = append(cmds, "-x")
		}
	}

	if len(cmds) == 0 {
		fmt.Println("No command specified")
		os.Exit(2)
	}

	if len(dirs) == 0 {
		var cwd, err = os.Getwd()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		dirs = []string{cwd}
	}

	fmt.Println("Watching", dirs, "for", cmds)

	watcher, err := fsnotify.NewWatcher()

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

	for _, dir := range dirs {
		subfolders := Subfolders(dir)
		for _, f := range subfolders {
			err = watcher.WatchFlags(f, fsnotify.FSN_CREATE|fsnotify.FSN_MODIFY)
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	var cmd *exec.Cmd

	for {
		select {
		case e := <-watcher.Event:
			matched, err := regexp.MatchString("\\.(go|c|h)$", e.Name)
			if err != nil {
				log.Println(err)
			}

			if !matched {
				continue
			}

			log.Println("Event:", e)

			if cmd != nil && cmd.ProcessState != nil && !cmd.ProcessState.Exited() {
				err := cmd.Process.Kill()
				if err != nil {
					log.Println(err)
				}
			}
			cmd = exec.Command(cmds[0], cmds[1:]...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Start()
		case err := <-watcher.Error:
			log.Println("Error:", err)
		}
	}

	watcher.Close()
}
<commit_msg>Check the file given in arguments until --<commit_after>package main

import "github.com/howeyc/fsnotify"
import "log"
import "flag"
import "fmt"
import "os"
import "os/exec"
import "regexp"

var versionStr = "0.1.0"

var goCommandSet = map[string][]string{
	"test":    []string{"go", "test"},
	"install": []string{"go", "install"},
	"build":   []string{"go", "build"},
}

func main() {
	var helpFlag = flag.Bool("h", false, "help")
	var buildFlag = flag.Bool("b", true, "build, default behavior")
	var testFlag = flag.Bool("t", false, "test")
	var installFlag = flag.Bool("i", false, "install")
	var versionFlag = flag.Bool("v", false, "version")
	var xFlag = flag.Bool("x", false, "show verbose command")

	flag.Parse()
	args := flag.Args()

	if *helpFlag {
		fmt.Println("Usage: gomon [options] [dir] [-- command]")
		fmt.Println("   -b build")
		fmt.Println("   -t test")
		fmt.Println("   -i install")
		fmt.Println("   -x show verbose command")
		fmt.Println("   -h help")
		os.Exit(0)
	}
	if *versionFlag {
		fmt.Printf("gomon %s\n", versionStr)
		os.Exit(0)
	}

	var dirs = []string{}
	var cmds = []string{}

	var takeDir = true
	for _, a := range args {
		var exists, _ = FileExists(a)
		if a == "--" || (takeDir && !exists) {
			takeDir = false
			continue
		}
		if takeDir {
			dirs = append(dirs, a)
		} else {
			cmds = append(cmds, a)
		}
	}

	if len(cmds) == 0 {
		if *testFlag {
			cmds = goCommandSet["test"]
		} else if *buildFlag {
			cmds = goCommandSet["build"]
		} else if *installFlag {
			cmds = goCommandSet["install"]
		}
		if *xFlag && len(cmds) > 0 {
			cmds = append(cmds, "-x")
		}
	}

	if len(cmds) == 0 {
		fmt.Println("No command specified")
		os.Exit(2)
	}

	if len(dirs) == 0 {
		var cwd, err = os.Getwd()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		dirs = []string{cwd}
	}

	fmt.Println("Watching", dirs, "for", cmds)

	watcher, err := fsnotify.NewWatcher()

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

	for _, dir := range dirs {
		subfolders := Subfolders(dir)
		for _, f := range subfolders {
			err = watcher.WatchFlags(f, fsnotify.FSN_CREATE|fsnotify.FSN_MODIFY)
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	var cmd *exec.Cmd

	for {
		select {
		case e := <-watcher.Event:
			matched, err := regexp.MatchString("\\.(go|c|h)$", e.Name)
			if err != nil {
				log.Println(err)
			}

			if !matched {
				continue
			}

			log.Println("Event:", e)

			if cmd != nil && cmd.ProcessState != nil && !cmd.ProcessState.Exited() {
				err := cmd.Process.Kill()
				if err != nil {
					log.Println(err)
				}
			}
			cmd = exec.Command(cmds[0], cmds[1:]...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Start()
		case err := <-watcher.Error:
			log.Println("Error:", err)
		}
	}

	watcher.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	_ "github.com/joho/godotenv/autoload"
	"github.com/urfave/cli"
)

var version string // build number set at compile-time

var (
	buildCommit string
)

func main() {
	app := cli.NewApp()
	app.Name = "slack"
	app.Usage = "slack plugin"
	app.Action = run
	app.Version = version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "webhook",
			Usage:  "slack webhook url",
			EnvVar: "SLACK_WEBHOOK,PLUGIN_WEBHOOK",
		},
		cli.StringFlag{
			Name:   "channel",
			Usage:  "slack channel",
			EnvVar: "PLUGIN_CHANNEL",
		},
		cli.StringFlag{
			Name:   "recipient",
			Usage:  "slack recipient",
			EnvVar: "PLUGIN_RECIPIENT",
		},
		cli.StringFlag{
			Name:   "username",
			Usage:  "slack username",
			Value:  "drone",
			EnvVar: "PLUGIN_USERNAME",
		},
		cli.StringFlag{
			Name:   "template",
			Usage:  "slack template",
			EnvVar: "PLUGIN_TEMPLATE",
		},
		cli.StringFlag{
			Name:   "image",
			Usage:  "slack image url",
			EnvVar: "PLUGIN_IMAGE_URL",
		},
		cli.StringFlag{
			Name:   "icon.url",
			Usage:  "slack icon url",
			EnvVar: "PLUGIN_ICON_URL",
		},
		cli.StringFlag{
			Name:   "icon.emoji",
			Usage:  "slack emoji url",
			EnvVar: "PLUGIN_ICON_EMOJI",
		},
		cli.StringFlag{
			Name:   "repo.owner",
			Usage:  "repository owner",
			EnvVar: "DRONE_REPO_OWNER",
		},
		cli.StringFlag{
			Name:   "repo.name",
			Usage:  "repository name",
			EnvVar: "DRONE_REPO_NAME",
		},
		cli.StringFlag{
			Name:   "commit.sha",
			Usage:  "git commit sha",
			EnvVar: "DRONE_COMMIT_SHA",
		},
		cli.StringFlag{
			Name:   "commit.branch",
			Value:  "master",
			Usage:  "git commit branch",
			EnvVar: "DRONE_COMMIT_BRANCH",
		},
		cli.StringFlag{
			Name:   "commit.author",
			Usage:  "git author name",
			EnvVar: "DRONE_COMMIT_AUTHOR",
		},
		cli.StringFlag{
			Name:   "build.event",
			Value:  "push",
			Usage:  "build event",
			EnvVar: "DRONE_BUILD_EVENT",
		},
		cli.IntFlag{
			Name:   "build.number",
			Usage:  "build number",
			EnvVar: "DRONE_BUILD_NUMBER",
		},
		cli.StringFlag{
			Name:   "build.status",
			Usage:  "build status",
			Value:  "success",
			EnvVar: "DRONE_BUILD_STATUS",
		},
		cli.StringFlag{
			Name:   "build.link",
			Usage:  "build link",
			EnvVar: "DRONE_BUILD_LINK",
		},
	}
	app.Run(os.Args)
}

func run(c *cli.Context) error {
	plugin := Plugin{
		Repo: Repo{
			Owner: c.String("repo.owner"),
			Name:  c.String("repo.name"),
		},
		Build: Build{
			Number: c.Int("build.number"),
			Event:  c.String("build.event"),
			Status: c.String("build.status"),
			Commit: c.String("commit.sha"),
			Branch: c.String("commit.branch"),
			Author: c.String("commit.author"),
			Link:   c.String("build.link"),
		},
		Config: Config{
			Webhook:   c.String("webhook"),
			Channel:   c.String("channel"),
			Recipient: c.String("recipient"),
			Username:  c.String("username"),
			Template:  c.String("template"),
			ImageURL:  c.String("image"),
			IconURL:   c.String("icon.url"),
			IconEmoji: c.String("icon.emoji"),
		},
	}

	return plugin.Exec()
}
<commit_msg>Added logrus, added env-file flag, disabled autoload<commit_after>package main

import (
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/joho/godotenv"
	"github.com/urfave/cli"
)

var version string // build number set at compile-time

func main() {
	app := cli.NewApp()
	app.Name = "slack plugin"
	app.Usage = "slack plugin"
	app.Action = run
	app.Version = version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "webhook",
			Usage:  "slack webhook url",
			EnvVar: "SLACK_WEBHOOK,PLUGIN_WEBHOOK",
		},
		cli.StringFlag{
			Name:   "channel",
			Usage:  "slack channel",
			EnvVar: "PLUGIN_CHANNEL",
		},
		cli.StringFlag{
			Name:   "recipient",
			Usage:  "slack recipient",
			EnvVar: "PLUGIN_RECIPIENT",
		},
		cli.StringFlag{
			Name:   "username",
			Usage:  "slack username",
			Value:  "drone",
			EnvVar: "PLUGIN_USERNAME",
		},
		cli.StringFlag{
			Name:   "template",
			Usage:  "slack template",
			EnvVar: "PLUGIN_TEMPLATE",
		},
		cli.StringFlag{
			Name:   "image",
			Usage:  "slack image url",
			EnvVar: "PLUGIN_IMAGE_URL",
		},
		cli.StringFlag{
			Name:   "icon.url",
			Usage:  "slack icon url",
			EnvVar: "PLUGIN_ICON_URL",
		},
		cli.StringFlag{
			Name:   "icon.emoji",
			Usage:  "slack emoji url",
			EnvVar: "PLUGIN_ICON_EMOJI",
		},
		cli.StringFlag{
			Name:   "repo.owner",
			Usage:  "repository owner",
			EnvVar: "DRONE_REPO_OWNER",
		},
		cli.StringFlag{
			Name:   "repo.name",
			Usage:  "repository name",
			EnvVar: "DRONE_REPO_NAME",
		},
		cli.StringFlag{
			Name:   "commit.sha",
			Usage:  "git commit sha",
			EnvVar: "DRONE_COMMIT_SHA",
		},
		cli.StringFlag{
			Name:   "commit.branch",
			Value:  "master",
			Usage:  "git commit branch",
			EnvVar: "DRONE_COMMIT_BRANCH",
		},
		cli.StringFlag{
			Name:   "commit.author",
			Usage:  "git author name",
			EnvVar: "DRONE_COMMIT_AUTHOR",
		},
		cli.StringFlag{
			Name:   "build.event",
			Value:  "push",
			Usage:  "build event",
			EnvVar: "DRONE_BUILD_EVENT",
		},
		cli.IntFlag{
			Name:   "build.number",
			Usage:  "build number",
			EnvVar: "DRONE_BUILD_NUMBER",
		},
		cli.StringFlag{
			Name:   "build.status",
			Usage:  "build status",
			Value:  "success",
			EnvVar: "DRONE_BUILD_STATUS",
		},
		cli.StringFlag{
			Name:   "build.link",
			Usage:  "build link",
			EnvVar: "DRONE_BUILD_LINK",
		},
		cli.StringFlag{
			Name:  "env-file",
			Usage: "source env file",
		},
	}

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

func run(c *cli.Context) error {
	if c.String("env-file") != "" {
		_ = godotenv.Load(c.String("env-file"))
	}

	plugin := Plugin{
		Repo: Repo{
			Owner: c.String("repo.owner"),
			Name:  c.String("repo.name"),
		},
		Build: Build{
			Number: c.Int("build.number"),
			Event:  c.String("build.event"),
			Status: c.String("build.status"),
			Commit: c.String("commit.sha"),
			Branch: c.String("commit.branch"),
			Author: c.String("commit.author"),
			Link:   c.String("build.link"),
		},
		Config: Config{
			Webhook:   c.String("webhook"),
			Channel:   c.String("channel"),
			Recipient: c.String("recipient"),
			Username:  c.String("username"),
			Template:  c.String("template"),
			ImageURL:  c.String("image"),
			IconURL:   c.String("icon.url"),
			IconEmoji: c.String("icon.emoji"),
		},
	}

	return plugin.Exec()
}
<|endoftext|>


<commit_before>package main

import (
	"collector"
	"encoding/json"
	"errors"
	"fmt"
	"helper"
	"log"
	"os"
	"os/signal"
	"sse"
	"time"
)

var (
	// The url that hosts the service that this program will interact with
	URL          = "http://mothership.serverstatusmonitoring.com"

	// The uri to be used to register the system this program will run on
	URIRegister  = "/register"

	// The uri where collected data will be sent to
	URICollector = "/collector"

	// The uri to check the upstatus of URL
	URIStatus    = "/status"

	// The hostname of the system this program will run on
	Hostname  = ""

	// The IP address of the system this program will run on
	IPAddress = ""

	// The file where we want to log event data and errors
	LogFile           = "/var/log/sphire-sse.log"

	// The configuration file we want to use
	ConfigurationFile = "/etc/sse/sse.conf"

	// The configuration instance (loads the above LogFile)
	Configuration     = new(Config)

	// Collect a snapshot and store in cache every X seconds
	CollectFrequencySeconds = 1

	// Report all snapshots in cache every Y seconds
	ReportFrequencySeconds  = 1

	// Instance of collector.CPU
	CPU     collector.CPU     = collector.CPU{}

	// Instance of collector.Disks
	Disks   collector.Disks   = collector.Disks{}

	// Instance of collector.Memory
	Memory  collector.Memory  = collector.Memory{}

	// Instance of collector.Network
	Network collector.Network = collector.Network{}

	// Instance of collector.System
	System  collector.System  = collector.System{}

	// The version of this program
	Version = "1.0.1"
)

/*
 Configuration struct is a direct map to the configuration located in the configuration JSON file.
*/
type Config struct {
	Identification struct {
		AccountID        string `json:"account_id"`
		OrganizationID   string `json:"organization_id"`
		OrganizationName string `json:"organization_name"`
		MachineNickname  string `json:"machine_nickname"`
	} `json:"identification"`
	Settings struct {
		Disk struct {
			IncludePartitionData bool `json:"include_partition_data"`
		} `json:"disk"`
		System struct {
			IncludeUsers bool `json:"include_users"`
		} `json:"system"`
		Reporting struct {
			CollectFrequencySeconds int `json:"collect_frequency_seconds"`
			ReportFrequencySeconds  int `json:"report_frequency_seconds"`
		} `json:"reporting"`
	} `json:"settings"`
}

func main() {
	// Define the global logger
	logger, err := os.OpenFile(LogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	HandleError(err)
	defer logger.Close()
	log.SetOutput(logger)

	// Load and parse configuration file
	file, err := os.Open(ConfigurationFile)
	HandleError(err)
	err = json.NewDecoder(file).Decode(Configuration)
	HandleError(err)

	// Set any parameters that need to be set
	CollectFrequencySeconds = Configuration.Settings.Reporting.CollectFrequencySeconds
	ReportFrequencySeconds = Configuration.Settings.Reporting.ReportFrequencySeconds

	var status helper.Status = helper.Status{}
	var status_result bool = status.CheckStatus(URL + URIStatus)
	if status_result == false {
		HandleError(errors.New("Mothership unreachable. Check your internet connection."))
	}

	// Perform system initialization
	var server_obj sse.Server = sse.Server{}
	server, ipAddress, hostname, version, error := server_obj.Initialize()
	HandleError(error)

	IPAddress = ipAddress
	Hostname = hostname
	Version = version

	// Perform registration
	body, err := sse.Register(map[string]interface{}{
		"configuration":     Configuration,
		"mothership_url":    URL,
		"register_uri":      URIRegister,
		"version":           Version,
		"collect_frequency": CollectFrequencySeconds,
		"report_frequency":  ReportFrequencySeconds,
		"hostname":          Hostname,
		"ip_address":        IPAddress,
		"log_file":          LogFile,
		"config_file":       ConfigurationFile,
	}, URL+URIRegister+"/"+Version)
	if err != nil {
		HandleError(errors.New("Unable to register this machine" + string(body)))
	}

	// Set up our collector
	var counter int = 0
	var snapshot sse.Snapshot = sse.Snapshot{}
	var cache sse.Cache = sse.Cache{
		AccountId:        Configuration.Identification.AccountID,
		OrganizationID:   Configuration.Identification.OrganizationID,
		OrganizationName: Configuration.Identification.OrganizationName,
		MachineNickname:  Configuration.Identification.MachineNickname,
		Version:          Version,
		Server:           server}

	ticker := time.NewTicker(time.Duration(CollectFrequencySeconds) * time.Second)
	death := make(chan os.Signal, 1)
	signal.Notify(death, os.Interrupt, os.Kill)

	go func() {
		for {
			select {
			case <-ticker.C : // send the updated time back via to the channel
				// reset the snapshot to an empty struct
				snapshot = sse.Snapshot{}

				// fill in the Snapshot struct and add to the cache
				cache.Node = append(cache.Node, snapshot.Collector(Configuration.Settings.Disk.IncludePartitionData,
						Configuration.Settings.System.IncludeUsers))
				counter++

				if counter > 0 && counter%ReportFrequencySeconds == 0 {
					cache.Sender(URL + URICollector)
					cache.Node = nil // Clear the Node Cache
					counter = 0
				}
			case <-death:
				fmt.Println("died")
				return
			}
		}
	}()

	return

}

func HandleError(err error) {
	if err != nil {
		log.Println(helper.Trace(err, "ERROR"))
		fmt.Println(err, "ERROR")
		os.Exit(1)
	}
}
<commit_msg>modify top level comments on exported vars in main.go<commit_after>package main

import (
	"collector"
	"encoding/json"
	"errors"
	"fmt"
	"helper"
	"log"
	"os"
	"os/signal"
	"sse"
	"time"
)

var (
	// URL is the url that hosts the service that this program will interact with
	URL          = "http://mothership.serverstatusmonitoring.com"

	// URIRegister is the uri to be used to register the system this program will run on
	URIRegister  = "/register"

	// URICollector is the uri where collected data will be sent to
	URICollector = "/collector"

	// URIStatus is the uri to check the upstatus of URL
	URIStatus    = "/status"

	// Hostname is the hostname of the system this program will run on
	Hostname  = ""

	// IPAddress is the IP address of the system this program will run on
	IPAddress = ""

	// LogFile is the file where we want to log event data and errors
	LogFile           = "/var/log/sphire-sse.log"

	// ConfigurationFile is the configuration file we want to use
	ConfigurationFile = "/etc/sse/sse.conf"

	// Configuration is the configuration instance (loads the above LogFile)
	Configuration     = new(Config)

	// Collect a snapshot and store in cache every X seconds
	CollectFrequencySeconds = 1

	// Report all snapshots in cache every Y seconds
	ReportFrequencySeconds  = 1

	// Instance of collector.CPU
	CPU     collector.CPU     = collector.CPU{}

	// Instance of collector.Disks
	Disks   collector.Disks   = collector.Disks{}

	// Instance of collector.Memory
	Memory  collector.Memory  = collector.Memory{}

	// Instance of collector.Network
	Network collector.Network = collector.Network{}

	// Instance of collector.System
	System  collector.System  = collector.System{}

	// The version of this program
	Version = "1.0.1"
)

/*
 Configuration struct is a direct map to the configuration located in the configuration JSON file.
*/
type Config struct {
	Identification struct {
		AccountID        string `json:"account_id"`
		OrganizationID   string `json:"organization_id"`
		OrganizationName string `json:"organization_name"`
		MachineNickname  string `json:"machine_nickname"`
	} `json:"identification"`
	Settings struct {
		Disk struct {
			IncludePartitionData bool `json:"include_partition_data"`
		} `json:"disk"`
		System struct {
			IncludeUsers bool `json:"include_users"`
		} `json:"system"`
		Reporting struct {
			CollectFrequencySeconds int `json:"collect_frequency_seconds"`
			ReportFrequencySeconds  int `json:"report_frequency_seconds"`
		} `json:"reporting"`
	} `json:"settings"`
}

func main() {
	// Define the global logger
	logger, err := os.OpenFile(LogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	HandleError(err)
	defer logger.Close()
	log.SetOutput(logger)

	// Load and parse configuration file
	file, err := os.Open(ConfigurationFile)
	HandleError(err)
	err = json.NewDecoder(file).Decode(Configuration)
	HandleError(err)

	// Set any parameters that need to be set
	CollectFrequencySeconds = Configuration.Settings.Reporting.CollectFrequencySeconds
	ReportFrequencySeconds = Configuration.Settings.Reporting.ReportFrequencySeconds

	var status helper.Status = helper.Status{}
	var status_result bool = status.CheckStatus(URL + URIStatus)
	if status_result == false {
		HandleError(errors.New("Mothership unreachable. Check your internet connection."))
	}

	// Perform system initialization
	var server_obj sse.Server = sse.Server{}
	server, ipAddress, hostname, version, error := server_obj.Initialize()
	HandleError(error)

	IPAddress = ipAddress
	Hostname = hostname
	Version = version

	// Perform registration
	body, err := sse.Register(map[string]interface{}{
		"configuration":     Configuration,
		"mothership_url":    URL,
		"register_uri":      URIRegister,
		"version":           Version,
		"collect_frequency": CollectFrequencySeconds,
		"report_frequency":  ReportFrequencySeconds,
		"hostname":          Hostname,
		"ip_address":        IPAddress,
		"log_file":          LogFile,
		"config_file":       ConfigurationFile,
	}, URL+URIRegister+"/"+Version)
	if err != nil {
		HandleError(errors.New("Unable to register this machine" + string(body)))
	}

	// Set up our collector
	var counter int = 0
	var snapshot sse.Snapshot = sse.Snapshot{}
	var cache sse.Cache = sse.Cache{
		AccountId:        Configuration.Identification.AccountID,
		OrganizationID:   Configuration.Identification.OrganizationID,
		OrganizationName: Configuration.Identification.OrganizationName,
		MachineNickname:  Configuration.Identification.MachineNickname,
		Version:          Version,
		Server:           server}

	ticker := time.NewTicker(time.Duration(CollectFrequencySeconds) * time.Second)
	death := make(chan os.Signal, 1)
	signal.Notify(death, os.Interrupt, os.Kill)

	go func() {
		for {
			select {
			case <-ticker.C : // send the updated time back via to the channel
				// reset the snapshot to an empty struct
				snapshot = sse.Snapshot{}

				// fill in the Snapshot struct and add to the cache
				cache.Node = append(cache.Node, snapshot.Collector(Configuration.Settings.Disk.IncludePartitionData,
						Configuration.Settings.System.IncludeUsers))
				counter++

				if counter > 0 && counter%ReportFrequencySeconds == 0 {
					cache.Sender(URL + URICollector)
					cache.Node = nil // Clear the Node Cache
					counter = 0
				}
			case <-death:
				fmt.Println("died")
				return
			}
		}
	}()

	return

}

func HandleError(err error) {
	if err != nil {
		log.Println(helper.Trace(err, "ERROR"))
		fmt.Println(err, "ERROR")
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package admin

import (
	"database/sql"
	"errors"
	"fmt"
	"github.com/extemporalgenome/slug"
	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
	"html/template"
	"io"
	"reflect"
	"strings"
)

var templates, _ = template.ParseGlob(
	"admin/templates/*.html",
)

type Admin struct {
	Router        *mux.Router
	Path          string
	Database      string
	Title         string
	NameTransform func(string) string

	Username string
	Password string
	sessions map[string]*session

	db          *sql.DB
	models      map[string]*model
	modelGroups []*modelGroup
}

// Setup registers page handlers and enables the admin.
func Setup(admin *Admin) (*Admin, error) {
	if len(admin.Title) == 0 {
		admin.Title = "Admin"
	}

	if len(admin.Username) == 0 || len(admin.Password) == 0 {
		return nil, errors.New("Username and/or password is missing")
	}

	admin.sessions = map[string]*session{}

	db, err := sql.Open("sqlite3", admin.Database)
	if err != nil {
		return nil, err
	}
	admin.db = db

	admin.models = map[string]*model{}
	admin.modelGroups = []*modelGroup{}

	sr := admin.Router.PathPrefix(admin.Path).Subrouter()
	sr.StrictSlash(true)
	sr.HandleFunc("/", admin.handlerWrapper(admin.handleIndex))
	sr.HandleFunc("/logout/", admin.handlerWrapper(admin.handleLogout))
	sr.HandleFunc("/model/{slug}/", admin.handlerWrapper(admin.handleList))
	sr.HandleFunc("/model/{slug}/new/", admin.handlerWrapper(admin.handleEdit))
	sr.HandleFunc("/model/{slug}/edit/{id}/", admin.handlerWrapper(admin.handleEdit))
	return admin, nil
}

// Group adds a model group to the admin front page.
// Use this to organize your models.
func (a *Admin) Group(name string) (*modelGroup, error) {
	if a.models == nil {
		return nil, errors.New("Must call .Serve() before adding groups and registering models")
	}

	group := &modelGroup{
		admin:  a,
		Name:   name,
		slug:   slug.SlugAscii(name),
		Models: []*model{},
	}

	a.modelGroups = append(a.modelGroups, group)

	return group, nil
}

type modelGroup struct {
	admin  *Admin
	Name   string
	slug   string
	Models []*model
}

type namedModel interface {
	AdminName() string
}

// RegisterModel adds a model to a model group.
func (g *modelGroup) RegisterModel(mdl interface{}) error {
	t := reflect.TypeOf(mdl)

	val := reflect.ValueOf(mdl)
	ind := reflect.Indirect(val)

	parts := strings.Split(t.String(), ".")
	name := parts[len(parts)-1]

	var tableName string
	if g.admin.NameTransform != nil {
		tableName = g.admin.NameTransform(name)
	} else {
		tableName = name
	}

	if named, ok := mdl.(namedModel); ok {
		name = named.AdminName()
	}

	am := model{
		Name:      name,
		Slug:      slug.SlugAscii(name),
		tableName: tableName,
		fields:    []Field{},
		instance:  mdl,
	}

	for i := 0; i < ind.NumField(); i++ {
		field := t.Elem().Field(i)
		fieldType := field.Type
		kind := fieldType.Kind()
		tag := field.Tag.Get("admin")
		if tag == "-" {
			continue
		}

		// Parse key=val / key options from struct tag, used for configuration later
		tagMap, err := parseTag(tag)
		if err != nil {
			panic(err)
		}

		// Expect pointers to be foreignkeys and foreignkeys to have the form Field[Id]
		fieldName := field.Name
		if kind == reflect.Ptr {
			fieldName += "Id"
		}

		// Transform struct keys to DB column names if needed
		var tableField string
		if g.admin.NameTransform != nil {
			tableField = g.admin.NameTransform(fieldName)
		} else {
			tableField = field.Name
		}

		// Choose Field
		var Field Field
		fmt.Println(kind)
		if FieldType, ok := tagMap["Field"]; ok {
			switch FieldType {
			case "url":
				Field = &URLField{BaseField: &BaseField{}}
			default:
				Field = &TextField{BaseField: &BaseField{}}
			}
		} else {
			switch kind {
			case reflect.String:
				Field = &TextField{BaseField: &BaseField{}}
			case reflect.Int, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				Field = &IntField{BaseField: &BaseField{}}
			case reflect.Float32, reflect.Float64:
				Field = &FloatField{BaseField: &BaseField{}}
			case reflect.Struct:
				Field = &TimeField{BaseField: &BaseField{}}
			default:
				fmt.Println("NOOO")
				Field = &TextField{BaseField: &BaseField{}}
			}
		}
		Field.Attrs().name = fieldName

		// Read relevant config options from the tagMap
		err = Field.Configure(tagMap)
		if err != nil {
			panic(err)
		}

		if label, ok := tagMap["label"]; ok {
			Field.Attrs().label = label
		} else {
			Field.Attrs().label = fieldName
		}

		Field.Attrs().columnName = tableField

		if _, ok := tagMap["list"]; ok {
			Field.Attrs().list = true
		}

		am.fields = append(am.fields, Field)
	}

	g.admin.models[am.Slug] = &am
	g.Models = append(g.Models, &am)
	return nil
}

type model struct {
	Name      string
	Slug      string
	fields    []Field
	tableName string
	instance  interface{}
}

func (m *model) renderForm(w io.Writer, data []interface{}, errors []string) {
	hasData := len(data) == len(m.fieldNames())
	var val interface{}
	for i, fieldName := range m.fieldNames() {
		if hasData {
			val = data[i]
		}
		var err string
		if errors != nil {
			err = errors[i]
		}
		field := m.fieldByName(fieldName)
		field.Render(w, val, err)
	}
}

func (m *model) fieldNames() []string {
	names := []string{}
	for _, field := range m.fields {
		names = append(names, field.Attrs().name)
	}
	return names
}

func (m *model) tableColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		names = append(names, field.Attrs().columnName)
	}
	return names
}

func (m *model) listColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		if !field.Attrs().list {
			continue
		}
		names = append(names, field.Attrs().label)
	}
	return names
}

func (m *model) listTableColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		if !field.Attrs().list {
			continue
		}
		names = append(names, field.Attrs().columnName)
	}
	return names
}

func (m *model) fieldByName(name string) Field {
	for _, field := range m.fields {
		if field.Attrs().name == name {
			return field
		}
	}
	return nil
}

func (a *Admin) modelURL(slug, action string) string {
	if _, ok := a.models[slug]; !ok {
		return a.Path
	}

	return fmt.Sprintf("%v/model/%v%v", a.Path, slug, action)
}
<commit_msg>Ugh, case insensitive rename...<commit_after>package admin

import (
	"database/sql"
	"errors"
	"fmt"
	"github.com/extemporalgenome/slug"
	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
	"html/template"
	"io"
	"reflect"
	"strings"
)

var templates, _ = template.ParseGlob(
	"admin/templates/*.html",
)

type Admin struct {
	Router        *mux.Router
	Path          string
	Database      string
	Title         string
	NameTransform func(string) string

	Username string
	Password string
	sessions map[string]*session

	db          *sql.DB
	models      map[string]*model
	modelGroups []*modelGroup
}

// Setup registers page handlers and enables the admin.
func Setup(admin *Admin) (*Admin, error) {
	if len(admin.Title) == 0 {
		admin.Title = "Admin"
	}

	if len(admin.Username) == 0 || len(admin.Password) == 0 {
		return nil, errors.New("Username and/or password is missing")
	}

	admin.sessions = map[string]*session{}

	db, err := sql.Open("sqlite3", admin.Database)
	if err != nil {
		return nil, err
	}
	admin.db = db

	admin.models = map[string]*model{}
	admin.modelGroups = []*modelGroup{}

	sr := admin.Router.PathPrefix(admin.Path).Subrouter()
	sr.StrictSlash(true)
	sr.HandleFunc("/", admin.handlerWrapper(admin.handleIndex))
	sr.HandleFunc("/logout/", admin.handlerWrapper(admin.handleLogout))
	sr.HandleFunc("/model/{slug}/", admin.handlerWrapper(admin.handleList))
	sr.HandleFunc("/model/{slug}/new/", admin.handlerWrapper(admin.handleEdit))
	sr.HandleFunc("/model/{slug}/edit/{id}/", admin.handlerWrapper(admin.handleEdit))
	return admin, nil
}

// Group adds a model group to the admin front page.
// Use this to organize your models.
func (a *Admin) Group(name string) (*modelGroup, error) {
	if a.models == nil {
		return nil, errors.New("Must call .Serve() before adding groups and registering models")
	}

	group := &modelGroup{
		admin:  a,
		Name:   name,
		slug:   slug.SlugAscii(name),
		Models: []*model{},
	}

	a.modelGroups = append(a.modelGroups, group)

	return group, nil
}

type modelGroup struct {
	admin  *Admin
	Name   string
	slug   string
	Models []*model
}

type namedModel interface {
	AdminName() string
}

// RegisterModel adds a model to a model group.
func (g *modelGroup) RegisterModel(mdl interface{}) error {
	t := reflect.TypeOf(mdl)

	val := reflect.ValueOf(mdl)
	ind := reflect.Indirect(val)

	parts := strings.Split(t.String(), ".")
	name := parts[len(parts)-1]

	var tableName string
	if g.admin.NameTransform != nil {
		tableName = g.admin.NameTransform(name)
	} else {
		tableName = name
	}

	if named, ok := mdl.(namedModel); ok {
		name = named.AdminName()
	}

	am := model{
		Name:      name,
		Slug:      slug.SlugAscii(name),
		tableName: tableName,
		fields:    []Field{},
		instance:  mdl,
	}

	for i := 0; i < ind.NumField(); i++ {
		refl := t.Elem().Field(i)
		fieldType := refl.Type
		kind := fieldType.Kind()
		tag := refl.Tag.Get("admin")
		if tag == "-" {
			continue
		}

		// Parse key=val / key options from struct tag, used for configuration later
		tagMap, err := parseTag(tag)
		if err != nil {
			panic(err)
		}

		// Expect pointers to be foreignkeys and foreignkeys to have the form Field[Id]
		fieldName := refl.Name
		if kind == reflect.Ptr {
			fieldName += "Id"
		}

		// Transform struct keys to DB column names if needed
		var tableField string
		if g.admin.NameTransform != nil {
			tableField = g.admin.NameTransform(fieldName)
		} else {
			tableField = refl.Name
		}

		// Choose Field
		var field Field
		fmt.Println(kind)
		if fieldType, ok := tagMap["Field"]; ok {
			switch fieldType {
			case "url":
				field = &URLField{BaseField: &BaseField{}}
			default:
				field = &TextField{BaseField: &BaseField{}}
			}
		} else {
			switch kind {
			case reflect.String:
				field = &TextField{BaseField: &BaseField{}}
			case reflect.Int, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				field = &IntField{BaseField: &BaseField{}}
			case reflect.Float32, reflect.Float64:
				field = &FloatField{BaseField: &BaseField{}}
			case reflect.Struct:
				field = &TimeField{BaseField: &BaseField{}}
			default:
				fmt.Println("NOOO")
				field = &TextField{BaseField: &BaseField{}}
			}
		}
		field.Attrs().name = fieldName

		// Read relevant config options from the tagMap
		err = field.Configure(tagMap)
		if err != nil {
			panic(err)
		}

		if label, ok := tagMap["label"]; ok {
			field.Attrs().label = label
		} else {
			field.Attrs().label = fieldName
		}

		field.Attrs().columnName = tableField

		if _, ok := tagMap["list"]; ok {
			field.Attrs().list = true
		}

		am.fields = append(am.fields, field)
	}

	g.admin.models[am.Slug] = &am
	g.Models = append(g.Models, &am)
	return nil
}

type model struct {
	Name      string
	Slug      string
	fields    []Field
	tableName string
	instance  interface{}
}

func (m *model) renderForm(w io.Writer, data []interface{}, errors []string) {
	hasData := len(data) == len(m.fieldNames())
	var val interface{}
	for i, fieldName := range m.fieldNames() {
		if hasData {
			val = data[i]
		}
		var err string
		if errors != nil {
			err = errors[i]
		}
		field := m.fieldByName(fieldName)
		field.Render(w, val, err)
	}
}

func (m *model) fieldNames() []string {
	names := []string{}
	for _, field := range m.fields {
		names = append(names, field.Attrs().name)
	}
	return names
}

func (m *model) tableColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		names = append(names, field.Attrs().columnName)
	}
	return names
}

func (m *model) listColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		if !field.Attrs().list {
			continue
		}
		names = append(names, field.Attrs().label)
	}
	return names
}

func (m *model) listTableColumns() []string {
	names := []string{}
	for _, field := range m.fields {
		if !field.Attrs().list {
			continue
		}
		names = append(names, field.Attrs().columnName)
	}
	return names
}

func (m *model) fieldByName(name string) Field {
	for _, field := range m.fields {
		if field.Attrs().name == name {
			return field
		}
	}
	return nil
}

func (a *Admin) modelURL(slug, action string) string {
	if _, ok := a.models[slug]; !ok {
		return a.Path
	}

	return fmt.Sprintf("%v/model/%v%v", a.Path, slug, action)
}
<|endoftext|>