<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/thomersch/grandine/lib/cugdf"
	"github.com/thomersch/grandine/lib/mvt"
	"github.com/thomersch/grandine/lib/spatial"
	"github.com/thomersch/grandine/lib/tile"
)

const (
	measureFilterTiming = false
	indexThreshold      = 100000000
)

type zmLvl []int

func (zm *zmLvl) String() string {
	return fmt.Sprintf("%d", *zm)
}

func (zm *zmLvl) Set(value string) error {
	v, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("%s (only integer values are allowed)", value)
	}
	*zm = append(*zm, v)
	return nil
}

var zoomlevels zmLvl

func main() {
	source := flag.String("src", "geo.geojson", "file to read from, supported formats: geojson, cugdf")
	target := flag.String("target", "tiles", "path where the tiles will be written")
	defaultLayer := flag.Bool("default-layer", true, "if no layer name is specified in the feature, whether it will be put into a default layer")
	workersNumber := flag.Int("workers", runtime.GOMAXPROCS(0), "number of workers")
	cpuProfile := flag.String("cpuprof", "", "writes CPU profiling data into a file")

	flag.Var(&zoomlevels, "zoom", "one or more zoom level of which the tiles will be rendered")
	flag.Parse()

	if len(zoomlevels) == 0 {
		log.Fatal("no zoom levels specified")
	}

	if len(*cpuProfile) != 0 {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	f, err := os.Open(*source)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	err = os.MkdirAll(*target, 0777)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("parsing input...")
	fc := spatial.FeatureCollection{}

	if strings.HasSuffix(strings.ToLower(*source), "geojson") {
		if err := json.NewDecoder(f).Decode(&fc); err != nil {
			log.Fatal(err)
		}
	} else {
		fc.Features, err = cugdf.Unmarshal(f)
		if err != nil {
			log.Fatal(err)
		}
	}

	if len(fc.Features) == 0 {
		log.Fatal("no features in input file")
	}

	log.Printf("read %d features", len(fc.Features))

	var bboxPts []spatial.Point
	for _, feat := range fc.Features {
		bb := feat.Geometry.BBox()
		bboxPts = append(bboxPts, bb.SW, bb.NE)
	}

	var tc []tile.ID
	for _, zoomlevel := range zoomlevels {
		tc = append(tc, tile.Coverage(spatial.Line(bboxPts).BBox(), zoomlevel)...)
	}
	log.Printf("attempting to generate %d tiles", len(tc))

	var fts spatial.Filterable
	if len(fc.Features)*len(tc) > indexThreshold {
		log.Println("building index...")
		fts = spatial.NewRTreeCollection(fc.Features...)
		log.Println("index complete")
	} else {
		fts = &fc
	}

	dtw := diskTileWriter{basedir: *target}
	dlm := defaultLayerMapper{defaultLayer: *defaultLayer}

	var (
		wg sync.WaitGroup
		ws = workerSlices(tc, *workersNumber)
	)
	for wrk := 0; wrk < len(ws); wrk++ {
		wg.Add(1)
		go func(i int) {
			generateTiles(ws[i], fts, &dtw, &dlm)
			wg.Done()
		}(wrk)
	}
	wg.Wait()
}

func workerSlices(tiles []tile.ID, wrkNum int) [][]tile.ID {
	var r [][]tile.ID
	if len(tiles) <= wrkNum {
		for t := 0; t < len(tiles); t++ {
			r = append(r, []tile.ID{tiles[t]})
		}
		return r
	}
	for wrkr := 0; wrkr < wrkNum; wrkr++ {
		start := (len(tiles) / wrkNum) * wrkr
		end := (len(tiles) / wrkNum) * (wrkr + 1)
		if wrkr == wrkNum-1 {
			end = len(tiles)
		}
		r = append(r, tiles[start:end])
	}
	return r
}

type diskTileWriter struct {
	basedir string
}

func (tw *diskTileWriter) WriteTile(tID tile.ID, buf []byte) error {
	err := os.MkdirAll(filepath.Join(tw.basedir, strconv.Itoa(tID.Z), strconv.Itoa(tID.X)), 0777)
	if err != nil {
		return err
	}
	tf, err := os.Create(filepath.Join(tw.basedir, strconv.Itoa(tID.Z), strconv.Itoa(tID.X), strconv.Itoa(tID.Y)+".mvt"))
	if err != nil {
		return err
	}
	defer tf.Close()
	_, err = tf.Write(buf)
	if err != nil {
		return err
	}
	return nil
}

type defaultLayerMapper struct {
	defaultLayer bool
}

func (dlm *defaultLayerMapper) LayerName(props map[string]interface{}) string {
	if layerName, ok := props["@layer"]; ok {
		return layerName.(string)
	}
	if dlm.defaultLayer {
		return "default"
	}
	return ""
}

type layerMapper interface {
	LayerName(map[string]interface{}) string
}

type tileWriter interface {
	WriteTile(tile.ID, []byte) error
}

func generateTiles(tIDs []tile.ID, features spatial.Filterable, tw tileWriter, lm layerMapper) {
	var (
		filterTime int64 //nans
		count      int64
		dur        time.Duration
		tStart     time.Time
	)
	for _, tID := range tIDs {
		// log.Printf("Generating %s", tID)
		var (
			layers = map[string][]spatial.Feature{}
			ln     string
		)
		tileClipBBox := tID.BBox()

		if measureFilterTiming {
			tStart = time.Now()
		}
		fts := features.Filter(tileClipBBox)
		if measureFilterTiming {
			dur = time.Now().Sub(tStart)
			filterTime += dur.Nanoseconds()
			count += 1
		}

		for _, feat := range fts {
			for _, geom := range feat.Geometry.ClipToBBox(tileClipBBox) {
				feat.Geometry = geom
				ln = lm.LayerName(feat.Props)
				if len(ln) != 0 {
					if _, ok := layers[ln]; !ok {
						layers[ln] = []spatial.Feature{feat}
					} else {
						layers[ln] = append(layers[ln], feat)
					}
				}
			}
		}
		if !anyFeatures(layers) {
			continue
		}
		buf, err := mvt.EncodeTile(layers, tID)
		if err != nil {
			log.Fatal(err)
		}

		err = tw.WriteTile(tID, buf)
		if err != nil {
			log.Fatal(err)
		}
	}
	if measureFilterTiming {
		log.Printf("avg duration: %v", time.Duration(filterTime/count))
	}
}

func anyFeatures(layers map[string][]spatial.Feature) bool {
	for _, ly := range layers {
		if len(ly) > 0 {
			return true
		}
	}
	return false
}
<commit_msg>cmd/tiler: remove index lookup timing calculation<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync"

	"github.com/thomersch/grandine/lib/cugdf"
	"github.com/thomersch/grandine/lib/mvt"
	"github.com/thomersch/grandine/lib/spatial"
	"github.com/thomersch/grandine/lib/tile"
)

const indexThreshold = 100000000

type zmLvl []int

func (zm *zmLvl) String() string {
	return fmt.Sprintf("%d", *zm)
}

func (zm *zmLvl) Set(value string) error {
	v, err := strconv.Atoi(value)
	if err != nil {
		return fmt.Errorf("%s (only integer values are allowed)", value)
	}
	*zm = append(*zm, v)
	return nil
}

var zoomlevels zmLvl

func main() {
	source := flag.String("src", "geo.geojson", "file to read from, supported formats: geojson, cugdf")
	target := flag.String("target", "tiles", "path where the tiles will be written")
	defaultLayer := flag.Bool("default-layer", true, "if no layer name is specified in the feature, whether it will be put into a default layer")
	workersNumber := flag.Int("workers", runtime.GOMAXPROCS(0), "number of workers")
	cpuProfile := flag.String("cpuprof", "", "writes CPU profiling data into a file")

	flag.Var(&zoomlevels, "zoom", "one or more zoom level of which the tiles will be rendered")
	flag.Parse()

	if len(zoomlevels) == 0 {
		log.Fatal("no zoom levels specified")
	}

	if len(*cpuProfile) != 0 {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	f, err := os.Open(*source)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	err = os.MkdirAll(*target, 0777)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("parsing input...")
	fc := spatial.FeatureCollection{}

	if strings.HasSuffix(strings.ToLower(*source), "geojson") {
		if err := json.NewDecoder(f).Decode(&fc); err != nil {
			log.Fatal(err)
		}
	} else {
		fc.Features, err = cugdf.Unmarshal(f)
		if err != nil {
			log.Fatal(err)
		}
	}

	if len(fc.Features) == 0 {
		log.Fatal("no features in input file")
	}

	log.Printf("read %d features", len(fc.Features))

	var bboxPts []spatial.Point
	for _, feat := range fc.Features {
		bb := feat.Geometry.BBox()
		bboxPts = append(bboxPts, bb.SW, bb.NE)
	}

	var tc []tile.ID
	for _, zoomlevel := range zoomlevels {
		tc = append(tc, tile.Coverage(spatial.Line(bboxPts).BBox(), zoomlevel)...)
	}
	log.Printf("attempting to generate %d tiles", len(tc))

	var fts spatial.Filterable
	if len(fc.Features)*len(tc) > indexThreshold {
		log.Println("building index...")
		fts = spatial.NewRTreeCollection(fc.Features...)
		log.Println("index complete")
	} else {
		fts = &fc
	}

	dtw := diskTileWriter{basedir: *target}
	dlm := defaultLayerMapper{defaultLayer: *defaultLayer}

	var (
		wg sync.WaitGroup
		ws = workerSlices(tc, *workersNumber)
	)
	for wrk := 0; wrk < len(ws); wrk++ {
		wg.Add(1)
		go func(i int) {
			generateTiles(ws[i], fts, &dtw, &dlm)
			wg.Done()
		}(wrk)
	}
	wg.Wait()
}

func workerSlices(tiles []tile.ID, wrkNum int) [][]tile.ID {
	var r [][]tile.ID
	if len(tiles) <= wrkNum {
		for t := 0; t < len(tiles); t++ {
			r = append(r, []tile.ID{tiles[t]})
		}
		return r
	}
	for wrkr := 0; wrkr < wrkNum; wrkr++ {
		start := (len(tiles) / wrkNum) * wrkr
		end := (len(tiles) / wrkNum) * (wrkr + 1)
		if wrkr == wrkNum-1 {
			end = len(tiles)
		}
		r = append(r, tiles[start:end])
	}
	return r
}

type diskTileWriter struct {
	basedir string
}

func (tw *diskTileWriter) WriteTile(tID tile.ID, buf []byte) error {
	err := os.MkdirAll(filepath.Join(tw.basedir, strconv.Itoa(tID.Z), strconv.Itoa(tID.X)), 0777)
	if err != nil {
		return err
	}
	tf, err := os.Create(filepath.Join(tw.basedir, strconv.Itoa(tID.Z), strconv.Itoa(tID.X), strconv.Itoa(tID.Y)+".mvt"))
	if err != nil {
		return err
	}
	defer tf.Close()
	_, err = tf.Write(buf)
	if err != nil {
		return err
	}
	return nil
}

type defaultLayerMapper struct {
	defaultLayer bool
}

func (dlm *defaultLayerMapper) LayerName(props map[string]interface{}) string {
	if layerName, ok := props["@layer"]; ok {
		return layerName.(string)
	}
	if dlm.defaultLayer {
		return "default"
	}
	return ""
}

type layerMapper interface {
	LayerName(map[string]interface{}) string
}

type tileWriter interface {
	WriteTile(tile.ID, []byte) error
}

func generateTiles(tIDs []tile.ID, features spatial.Filterable, tw tileWriter, lm layerMapper) {
	for _, tID := range tIDs {
		// log.Printf("Generating %s", tID)
		var (
			layers = map[string][]spatial.Feature{}
			ln     string
		)
		tileClipBBox := tID.BBox()

		for _, feat := range features.Filter(tileClipBBox) {
			for _, geom := range feat.Geometry.ClipToBBox(tileClipBBox) {
				feat.Geometry = geom
				ln = lm.LayerName(feat.Props)
				if len(ln) != 0 {
					if _, ok := layers[ln]; !ok {
						layers[ln] = []spatial.Feature{feat}
					} else {
						layers[ln] = append(layers[ln], feat)
					}
				}
			}
		}
		if !anyFeatures(layers) {
			continue
		}
		buf, err := mvt.EncodeTile(layers, tID)
		if err != nil {
			log.Fatal(err)
		}

		err = tw.WriteTile(tID, buf)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func anyFeatures(layers map[string][]spatial.Feature) bool {
	for _, ly := range layers {
		if len(ly) > 0 {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/coreos/pkg/capnslog"
	"github.com/dustin/go-humanize"
	"github.com/ricochet2200/go-disk-usage/du"
	"github.com/spf13/cobra"

	"github.com/coreos/torus"
	"github.com/coreos/torus/blockset"
	"github.com/coreos/torus/distributor"
	"github.com/coreos/torus/internal/flagconfig"
	"github.com/coreos/torus/internal/http"
	"github.com/coreos/torus/models"
	"github.com/coreos/torus/ring"

	// Register all the possible drivers.
	_ "github.com/coreos/torus/block"
	_ "github.com/coreos/torus/metadata/etcd"
	_ "github.com/coreos/torus/metadata/temp"
	_ "github.com/coreos/torus/storage"
)

var (
	dataDir     string
	httpAddress string
	peerAddress string
	sizeStr     string
	size        uint64
	host        string
	port        int
	debugInit   bool
	autojoin    bool
	logpkg      string
	cfg         torus.Config

	debug   bool
	version bool
)

var rootCommand = &cobra.Command{
	Use:    "torusd",
	Short:  "Torus distributed storage",
	Long:   `The torus distributed storage server.`,
	PreRun: configureServer,
	Run:    runServer,
}

func init() {
	rootCommand.PersistentFlags().StringVarP(&dataDir, "data-dir", "", "torus-data", "Path to the data directory")
	rootCommand.PersistentFlags().BoolVarP(&debug, "debug", "", false, "Turn on debug output")
	rootCommand.PersistentFlags().BoolVarP(&debugInit, "debug-init", "", false, "Run a default init for the MDS if one doesn't exist")
	rootCommand.PersistentFlags().StringVarP(&host, "host", "", "", "Host to listen on for HTTP")
	rootCommand.PersistentFlags().IntVarP(&port, "port", "", 4321, "Port to listen on for HTTP")
	rootCommand.PersistentFlags().StringVarP(&peerAddress, "peer-address", "", "", "Address to listen on for intra-cluster data")
	rootCommand.PersistentFlags().StringVarP(&sizeStr, "size", "", "1GiB", "How much disk space to use for this storage node")
	rootCommand.PersistentFlags().StringVarP(&logpkg, "logpkg", "", "", "Specific package logging")
	rootCommand.PersistentFlags().BoolVarP(&autojoin, "auto-join", "", false, "Automatically join the storage pool")
	rootCommand.PersistentFlags().BoolVarP(&version, "version", "", false, "Print version info and exit")
	flagconfig.AddConfigFlags(rootCommand.PersistentFlags())
}

func main() {
	if err := rootCommand.Execute(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func configureServer(cmd *cobra.Command, args []string) {
	if version {
		fmt.Printf("torusd\nVersion: %s\n", torus.Version)
		os.Exit(0)
	}
	switch {
	case debug:
		capnslog.SetGlobalLogLevel(capnslog.DEBUG)
	default:
		capnslog.SetGlobalLogLevel(capnslog.INFO)
	}
	if logpkg != "" {
		capnslog.SetGlobalLogLevel(capnslog.NOTICE)
		rl := capnslog.MustRepoLogger("github.com/coreos/torus")
		llc, err := rl.ParseLogLevelConfig(logpkg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing logpkg: %s\n", err)
			os.Exit(1)
		}
		rl.SetLogLevel(llc)
	}

	if host != "" {
		httpAddress = fmt.Sprintf("%s:%d", host, port)
	}

	var err error
	if strings.Contains(sizeStr, "%") {
		percent, err := parsePercentage(sizeStr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing size %s: %s\n", sizeStr, err)
			os.Exit(1)
		}
		directory, _ := filepath.Abs(dataDir)
		size = du.NewDiskUsage(directory).Size() * percent / 100
	} else {
		size, err = humanize.ParseBytes(sizeStr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing size %s: %s\n", sizeStr, err)
			os.Exit(1)
		}
	}

	cfg = flagconfig.BuildConfigFromFlags()
	cfg.DataDir = dataDir
	cfg.StorageSize = size
}

func parsePercentage(percentString string) (uint64, error) {
	sizePercent := strings.Split(percentString, "%")[0]
	sizeNumber, err := strconv.Atoi(sizePercent)
	if err != nil {
		return 0, err
	}
	if sizeNumber < 1 || sizeNumber > 100 {
		return 0, errors.New(fmt.Sprintf("invalid size %d; must be between 1%% and 100%%", sizeNumber))
	}
	return uint64(sizeNumber), nil
}

func runServer(cmd *cobra.Command, args []string) {

	var (
		srv *torus.Server
		err error
	)
	switch {
	case cfg.MetadataAddress == "":
		srv, err = torus.NewServer(cfg, "temp", "mfile")
	case debugInit:
		err = torus.InitMDS("etcd", cfg, torus.GlobalMetadata{
			BlockSize:        512 * 1024,
			DefaultBlockSpec: blockset.MustParseBlockLayerSpec("crc,base"),
			INodeReplication: 2,
		}, ring.Ketama)
		if err != nil {
			if err == torus.ErrExists {
				fmt.Println("debug-init: Already exists")
			} else {
				fmt.Printf("Couldn't debug-init: %s\n", err)
				os.Exit(1)
			}
		}
		fallthrough
	default:
		srv, err = torus.NewServer(cfg, "etcd", "mfile")
	}
	if err != nil {
		fmt.Printf("Couldn't start: %s\n", err)
		os.Exit(1)
	}

	if autojoin {
		err = doAutojoin(srv)
		if err != nil {
			fmt.Printf("Couldn't auto-join: %s\n", err)
			os.Exit(1)
		}
	}

	mainClose := make(chan bool)
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	if peerAddress != "" {
		var u *url.URL

		u, err = url.Parse(peerAddress)
		if err != nil {
			fmt.Printf("Couldn't parse peer address %s: %s\n", peerAddress, err)
			os.Exit(1)
		}

		if u.Scheme == "" {
			fmt.Printf("Peer address %s does not have URL scheme (http:// or tdp://)\n", peerAddress)
			os.Exit(1)
		}

		err = distributor.ListenReplication(srv, u)
	} else {
		err = distributor.OpenReplication(srv)
	}

	defer srv.Close()
	go func() {
		for _ = range signalChan {
			fmt.Println("\nReceived an interrupt, stopping services...")
			close(mainClose)
			os.Exit(0)
		}
	}()

	if err != nil {
		fmt.Println("couldn't use server:", err)
		os.Exit(1)
	}
	if httpAddress != "" {
		http.ServeHTTP(httpAddress, srv)
	}
	// Wait
	<-mainClose
}

func doAutojoin(s *torus.Server) error {
	for {
		ring, err := s.MDS.GetRing()
		if err != nil {
			fmt.Fprintf(os.Stderr, "couldn't get ring: %v\n", err)
			return err
		}
		var newRing torus.Ring
		if r, ok := ring.(torus.RingAdder); ok {
			newRing, err = r.AddPeers(torus.PeerInfoList{
				&models.PeerInfo{
					UUID:        s.MDS.UUID(),
					TotalBlocks: s.Blocks.NumBlocks(),
				},
			})
		} else {
			fmt.Fprintf(os.Stderr, "current ring type cannot support auto-adding\n")
			return err
		}
		if err == torus.ErrExists {
			// We're already a member; we're coming back up.
			return nil
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "couldn't add peer to ring: %v", err)
			return err
		}
		err = s.MDS.SetRing(newRing)
		if err == torus.ErrNonSequentialRing {
			continue
		}
		return err
	}
}
<commit_msg>torusd: support completion option to generate bash-completion shell<commit_after>package main

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/coreos/pkg/capnslog"
	"github.com/dustin/go-humanize"
	"github.com/ricochet2200/go-disk-usage/du"
	"github.com/spf13/cobra"

	"github.com/coreos/torus"
	"github.com/coreos/torus/blockset"
	"github.com/coreos/torus/distributor"
	"github.com/coreos/torus/internal/flagconfig"
	"github.com/coreos/torus/internal/http"
	"github.com/coreos/torus/models"
	"github.com/coreos/torus/ring"

	// Register all the possible drivers.
	_ "github.com/coreos/torus/block"
	_ "github.com/coreos/torus/metadata/etcd"
	_ "github.com/coreos/torus/metadata/temp"
	_ "github.com/coreos/torus/storage"
)

var (
	dataDir     string
	httpAddress string
	peerAddress string
	sizeStr     string
	size        uint64
	host        string
	port        int
	debugInit   bool
	autojoin    bool
	logpkg      string
	cfg         torus.Config

	debug      bool
	version    bool
	completion bool
)

var rootCommand = &cobra.Command{
	Use:    "torusd",
	Short:  "Torus distributed storage",
	Long:   `The torus distributed storage server.`,
	PreRun: configureServer,
	Run:    runServer,
}

func init() {
	rootCommand.PersistentFlags().StringVarP(&dataDir, "data-dir", "", "torus-data", "Path to the data directory")
	rootCommand.PersistentFlags().BoolVarP(&debug, "debug", "", false, "Turn on debug output")
	rootCommand.PersistentFlags().BoolVarP(&debugInit, "debug-init", "", false, "Run a default init for the MDS if one doesn't exist")
	rootCommand.PersistentFlags().StringVarP(&host, "host", "", "", "Host to listen on for HTTP")
	rootCommand.PersistentFlags().IntVarP(&port, "port", "", 4321, "Port to listen on for HTTP")
	rootCommand.PersistentFlags().StringVarP(&peerAddress, "peer-address", "", "", "Address to listen on for intra-cluster data")
	rootCommand.PersistentFlags().StringVarP(&sizeStr, "size", "", "1GiB", "How much disk space to use for this storage node")
	rootCommand.PersistentFlags().StringVarP(&logpkg, "logpkg", "", "", "Specific package logging")
	rootCommand.PersistentFlags().BoolVarP(&autojoin, "auto-join", "", false, "Automatically join the storage pool")
	rootCommand.PersistentFlags().BoolVarP(&version, "version", "", false, "Print version info and exit")
	rootCommand.PersistentFlags().BoolVarP(&completion, "completion", "", false, "Output bash completion code")
	flagconfig.AddConfigFlags(rootCommand.PersistentFlags())
}

func main() {
	if err := rootCommand.Execute(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func configureServer(cmd *cobra.Command, args []string) {
	if version {
		fmt.Printf("torusd\nVersion: %s\n", torus.Version)
		os.Exit(0)
	}
	switch {
	case debug:
		capnslog.SetGlobalLogLevel(capnslog.DEBUG)
	default:
		capnslog.SetGlobalLogLevel(capnslog.INFO)
	}
	if logpkg != "" {
		capnslog.SetGlobalLogLevel(capnslog.NOTICE)
		rl := capnslog.MustRepoLogger("github.com/coreos/torus")
		llc, err := rl.ParseLogLevelConfig(logpkg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing logpkg: %s\n", err)
			os.Exit(1)
		}
		rl.SetLogLevel(llc)
	}

	if host != "" {
		httpAddress = fmt.Sprintf("%s:%d", host, port)
	}

	var err error
	if strings.Contains(sizeStr, "%") {
		percent, err := parsePercentage(sizeStr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing size %s: %s\n", sizeStr, err)
			os.Exit(1)
		}
		directory, _ := filepath.Abs(dataDir)
		size = du.NewDiskUsage(directory).Size() * percent / 100
	} else {
		size, err = humanize.ParseBytes(sizeStr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error parsing size %s: %s\n", sizeStr, err)
			os.Exit(1)
		}
	}

	cfg = flagconfig.BuildConfigFromFlags()
	cfg.DataDir = dataDir
	cfg.StorageSize = size
}

func parsePercentage(percentString string) (uint64, error) {
	sizePercent := strings.Split(percentString, "%")[0]
	sizeNumber, err := strconv.Atoi(sizePercent)
	if err != nil {
		return 0, err
	}
	if sizeNumber < 1 || sizeNumber > 100 {
		return 0, errors.New(fmt.Sprintf("invalid size %d; must be between 1%% and 100%%", sizeNumber))
	}
	return uint64(sizeNumber), nil
}

func runServer(cmd *cobra.Command, args []string) {
	if completion {
		cmd.Root().GenBashCompletion(os.Stdout)
		os.Exit(0)
	}

	var (
		srv *torus.Server
		err error
	)
	switch {
	case cfg.MetadataAddress == "":
		srv, err = torus.NewServer(cfg, "temp", "mfile")
	case debugInit:
		err = torus.InitMDS("etcd", cfg, torus.GlobalMetadata{
			BlockSize:        512 * 1024,
			DefaultBlockSpec: blockset.MustParseBlockLayerSpec("crc,base"),
			INodeReplication: 2,
		}, ring.Ketama)
		if err != nil {
			if err == torus.ErrExists {
				fmt.Println("debug-init: Already exists")
			} else {
				fmt.Printf("Couldn't debug-init: %s\n", err)
				os.Exit(1)
			}
		}
		fallthrough
	default:
		srv, err = torus.NewServer(cfg, "etcd", "mfile")
	}
	if err != nil {
		fmt.Printf("Couldn't start: %s\n", err)
		os.Exit(1)
	}

	if autojoin {
		err = doAutojoin(srv)
		if err != nil {
			fmt.Printf("Couldn't auto-join: %s\n", err)
			os.Exit(1)
		}
	}

	mainClose := make(chan bool)
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	if peerAddress != "" {
		var u *url.URL

		u, err = url.Parse(peerAddress)
		if err != nil {
			fmt.Printf("Couldn't parse peer address %s: %s\n", peerAddress, err)
			os.Exit(1)
		}

		if u.Scheme == "" {
			fmt.Printf("Peer address %s does not have URL scheme (http:// or tdp://)\n", peerAddress)
			os.Exit(1)
		}

		err = distributor.ListenReplication(srv, u)
	} else {
		err = distributor.OpenReplication(srv)
	}

	defer srv.Close()
	go func() {
		for _ = range signalChan {
			fmt.Println("\nReceived an interrupt, stopping services...")
			close(mainClose)
			os.Exit(0)
		}
	}()

	if err != nil {
		fmt.Println("couldn't use server:", err)
		os.Exit(1)
	}
	if httpAddress != "" {
		http.ServeHTTP(httpAddress, srv)
	}
	// Wait
	<-mainClose
}

func doAutojoin(s *torus.Server) error {
	for {
		ring, err := s.MDS.GetRing()
		if err != nil {
			fmt.Fprintf(os.Stderr, "couldn't get ring: %v\n", err)
			return err
		}
		var newRing torus.Ring
		if r, ok := ring.(torus.RingAdder); ok {
			newRing, err = r.AddPeers(torus.PeerInfoList{
				&models.PeerInfo{
					UUID:        s.MDS.UUID(),
					TotalBlocks: s.Blocks.NumBlocks(),
				},
			})
		} else {
			fmt.Fprintf(os.Stderr, "current ring type cannot support auto-adding\n")
			return err
		}
		if err == torus.ErrExists {
			// We're already a member; we're coming back up.
			return nil
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "couldn't add peer to ring: %v", err)
			return err
		}
		err = s.MDS.SetRing(newRing)
		if err == torus.ErrNonSequentialRing {
			continue
		}
		return err
	}
}
<|endoftext|>


<commit_before>package torrent

import (
	"errors"
	"io"
	"os"
	"sync"
)

// Accesses torrent data via a client.
type Reader struct {
	t          *Torrent
	responsive bool
	// Ensure operations that change the position are exclusive, like Read()
	// and Seek().
	opMu sync.Mutex

	// Required when modifying pos and readahead, or reading them without
	// opMu.
	mu        sync.Mutex
	pos       int64
	readahead int64
}

var _ io.ReadCloser = &Reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *Reader) SetResponsive() {
	r.responsive = true
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *Reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.mu.Unlock()
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	r.tickleClient()
}

func (r *Reader) readable(off int64) (ret bool) {
	if r.torrentClosed() {
		return true
	}
	req, ok := r.t.torrent.offsetRequest(off)
	if !ok {
		panic(off)
	}
	if r.responsive {
		return r.t.torrent.haveChunk(req)
	}
	return r.t.torrent.pieceComplete(int(req.Index))
}

// How many bytes are available to read. Max is the most we could require.
func (r *Reader) available(off, max int64) (ret int64) {
	for max > 0 {
		req, ok := r.t.torrent.offsetRequest(off)
		if !ok {
			break
		}
		if !r.t.torrent.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.torrent.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *Reader) tickleClient() {
	r.t.torrent.readersChanged()
}

func (r *Reader) waitReadable(off int64) {
	// We may have been sent back here because we were told we could read but
	// it failed.
	r.tickleClient()
	r.t.cl.event.Wait()
}

func (r *Reader) Read(b []byte) (n int, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	for len(b) != 0 {
		var n1 int
		n1, err = r.readOnceAt(b, r.pos)
		if n1 == 0 {
			if err == nil {
				panic("expected error")
			}
			break
		}
		b = b[n1:]
		n += n1
		r.mu.Lock()
		r.pos += int64(n1)
		r.mu.Unlock()
	}
	if r.pos >= r.t.torrent.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

// Safe to call with or without client lock.
func (r *Reader) torrentClosed() bool {
	return r.t.torrent.isClosed()
}

// Wait until some data should be available to read. Tickles the client if it
// isn't. Returns how much should be readable without blocking.
func (r *Reader) waitAvailable(pos, wanted int64) (avail int64) {
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	for !r.readable(pos) {
		r.waitReadable(pos)
	}
	return r.available(pos, wanted)
}

// Performs at most one successful read to torrent storage.
func (r *Reader) readOnceAt(b []byte, pos int64) (n int, err error) {
	if pos >= r.t.torrent.length {
		err = io.EOF
		return
	}
	for {
		avail := r.waitAvailable(pos, int64(len(b)))
		if avail == 0 {
			if r.torrentClosed() {
				err = errors.New("torrent closed")
				return
			}
		}
		b1 := b[:avail]
		pi := int(pos / r.t.Info().PieceLength)
		ip := r.t.Info().Piece(pi)
		po := pos % ip.Length()
		if int64(len(b1)) > ip.Length()-po {
			b1 = b1[:ip.Length()-po]
		}
		n, err = r.t.torrent.readAt(b1, pos)
		if n != 0 {
			return
		}
		// log.Printf("%s: error reading from torrent storage pos=%d: %s", r.t, pos, err)
		r.t.cl.mu.Lock()
		r.t.torrent.updateAllPieceCompletions()
		r.t.torrent.updatePiecePriorities()
		r.t.cl.mu.Unlock()
	}
}

func (r *Reader) Close() error {
	r.t.deleteReader(r)
	r.t = nil
	return nil
}

func (r *Reader) posChanged() {
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	r.t.torrent.readersChanged()
}

func (r *Reader) Seek(off int64, whence int) (ret int64, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()

	r.mu.Lock()
	switch whence {
	case os.SEEK_SET:
		r.pos = off
	case os.SEEK_CUR:
		r.pos += off
	case os.SEEK_END:
		r.pos = r.t.torrent.Info.TotalLength() + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos
	r.mu.Unlock()

	r.posChanged()
	return
}

func (r *Reader) Torrent() *Torrent {
	return r.t
}
<commit_msg>Suppress piece read errors when data is obtained<commit_after>package torrent

import (
	"errors"
	"io"
	"os"
	"sync"

	"github.com/anacrolix/missinggo"
)

// Accesses torrent data via a client.
type Reader struct {
	t          *Torrent
	responsive bool
	// Ensure operations that change the position are exclusive, like Read()
	// and Seek().
	opMu sync.Mutex

	// Required when modifying pos and readahead, or reading them without
	// opMu.
	mu        sync.Mutex
	pos       int64
	readahead int64
}

var _ io.ReadCloser = &Reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *Reader) SetResponsive() {
	r.responsive = true
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *Reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.mu.Unlock()
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	r.tickleClient()
}

func (r *Reader) readable(off int64) (ret bool) {
	if r.torrentClosed() {
		return true
	}
	req, ok := r.t.torrent.offsetRequest(off)
	if !ok {
		panic(off)
	}
	if r.responsive {
		return r.t.torrent.haveChunk(req)
	}
	return r.t.torrent.pieceComplete(int(req.Index))
}

// How many bytes are available to read. Max is the most we could require.
func (r *Reader) available(off, max int64) (ret int64) {
	for max > 0 {
		req, ok := r.t.torrent.offsetRequest(off)
		if !ok {
			break
		}
		if !r.t.torrent.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.torrent.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *Reader) tickleClient() {
	r.t.torrent.readersChanged()
}

func (r *Reader) waitReadable(off int64) {
	// We may have been sent back here because we were told we could read but
	// it failed.
	r.tickleClient()
	r.t.cl.event.Wait()
}

func (r *Reader) Read(b []byte) (n int, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	for len(b) != 0 {
		var n1 int
		n1, err = r.readOnceAt(b, r.pos)
		if n1 == 0 {
			if err == nil {
				panic("expected error")
			}
			break
		}
		b = b[n1:]
		n += n1
		r.mu.Lock()
		r.pos += int64(n1)
		r.mu.Unlock()
	}
	if r.pos >= r.t.torrent.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

// Safe to call with or without client lock.
func (r *Reader) torrentClosed() bool {
	return r.t.torrent.isClosed()
}

// Wait until some data should be available to read. Tickles the client if it
// isn't. Returns how much should be readable without blocking.
func (r *Reader) waitAvailable(pos, wanted int64) (avail int64) {
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	for !r.readable(pos) {
		r.waitReadable(pos)
	}
	return r.available(pos, wanted)
}

// Performs at most one successful read to torrent storage.
func (r *Reader) readOnceAt(b []byte, pos int64) (n int, err error) {
	if pos >= r.t.torrent.length {
		err = io.EOF
		return
	}
	for {
		avail := r.waitAvailable(pos, int64(len(b)))
		if avail == 0 {
			if r.torrentClosed() {
				err = errors.New("torrent closed")
				return
			}
		}
		b1 := b[:avail]
		pi := int(pos / r.t.Info().PieceLength)
		ip := r.t.Info().Piece(pi)
		po := pos % ip.Length()
		missinggo.LimitLen(&b1, ip.Length()-po)
		n, err = r.t.torrent.readAt(b1, pos)
		if n != 0 {
			err = nil
			return
		}
		// log.Printf("%s: error reading from torrent storage pos=%d: %s", r.t, pos, err)
		r.t.cl.mu.Lock()
		r.t.torrent.updateAllPieceCompletions()
		r.t.torrent.updatePiecePriorities()
		r.t.cl.mu.Unlock()
	}
}

func (r *Reader) Close() error {
	r.t.deleteReader(r)
	r.t = nil
	return nil
}

func (r *Reader) posChanged() {
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	r.t.torrent.readersChanged()
}

func (r *Reader) Seek(off int64, whence int) (ret int64, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()

	r.mu.Lock()
	switch whence {
	case os.SEEK_SET:
		r.pos = off
	case os.SEEK_CUR:
		r.pos += off
	case os.SEEK_END:
		r.pos = r.t.torrent.Info.TotalLength() + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos
	r.mu.Unlock()

	r.posChanged()
	return
}

func (r *Reader) Torrent() *Torrent {
	return r.t
}
<|endoftext|>


<commit_before>package disruptor

import "time"

type Reader struct {
	read     *Cursor
	written  *Cursor
	upstream Barrier
	consumer Consumer
	ready    bool
} // TODO: padding???

func NewReader(read, written *Cursor, upstream Barrier, consumer Consumer) *Reader {
	return &Reader{
		read:     read,
		written:  written,
		upstream: upstream,
		consumer: consumer,
		ready:    false,
	}
}

func (this *Reader) Start() {
	this.ready = true
	go this.receive()
}
func (this *Reader) Stop() {
	this.ready = false
}

func (this *Reader) receive() {
	previous := this.read.Sequence
	idling, gating := 0, 0

	for {
		lower := previous + 1
		upper := this.upstream.Read(lower)

		if lower <= upper {
			this.consumer.Consume(lower, upper)
			this.read.Sequence = upper
			previous = upper
		} else if upper = this.written.Load(); lower <= upper {
			// Gating--TODO: wait strategy (provide gating count to wait strategy for phased backoff)
			gating++
			idling = 0
		} else if this.ready {
			// Idling--TODO: wait strategy (provide idling count to wait strategy for phased backoff)
			idling++
			gating = 0
		} else {
			break
		}

		// sleeping increases the batch size which reduces the number of read commits
		// which drastically reduces the cost; longer sleeps = larger batches = less expensive commits
		time.Sleep(time.Microsecond)
	}
}
<commit_msg>Added TODOs to reader.<commit_after>package disruptor

import "time"

type Reader struct {
	read     *Cursor
	written  *Cursor
	upstream Barrier
	consumer Consumer
	ready    bool
} // TODO: padding???

func NewReader(read, written *Cursor, upstream Barrier, consumer Consumer) *Reader {
	return &Reader{
		read:     read,
		written:  written,
		upstream: upstream,
		consumer: consumer,
		ready:    false,
	}
}

func (this *Reader) Start() {
	this.ready = true
	go this.receive()
}
func (this *Reader) Stop() {
	this.ready = false
}

func (this *Reader) receive() {
	previous := this.read.Sequence // TODO: this.read.Load()
	idling, gating := 0, 0

	for {
		lower := previous + 1
		upper := this.upstream.Read(lower)

		if lower <= upper {
			this.consumer.Consume(lower, upper)
			this.read.Sequence = upper // TODO: this.read.Commit()
			previous = upper
		} else if upper = this.written.Load(); lower <= upper {
			// Gating--TODO: wait strategy (provide gating count to wait strategy for phased backoff)
			gating++
			idling = 0
		} else if this.ready {
			// Idling--TODO: wait strategy (provide idling count to wait strategy for phased backoff)
			idling++
			gating = 0
		} else {
			break
		}

		// sleeping increases the batch size which reduces the number of read commits
		// which drastically reduces the cost; longer sleeps = larger batches = less expensive commits
		time.Sleep(time.Microsecond)
	}
}
<|endoftext|>


<commit_before>package incus

import "log"

type RedisConsumer struct {
	commands <-chan RedisCommand
	pool     *redisPool
}

func NewRedisConsumer(commands <-chan RedisCommand, pool *redisPool) *RedisConsumer {
	consumer := &RedisConsumer{
		commands: commands,
		pool:     pool,
	}

	go consumer.ConsumeForever()

	return consumer
}

func (r *RedisConsumer) ConsumeForever() {
	for {
		command := <-r.commands

		if DEBUG {
			log.Println("Dequeued one command in consumer")
		}

		conn, success := r.pool.Get()

		if success {
			command(conn)
		} else {
			log.Println("Failed to get redis connection")
		}
	}
}
<commit_msg>fix: Return the conn to the pool after consuming<commit_after>package incus

import "log"

type RedisConsumer struct {
	commands <-chan RedisCommand
	pool     *redisPool
}

func NewRedisConsumer(commands <-chan RedisCommand, pool *redisPool) *RedisConsumer {
	consumer := &RedisConsumer{
		commands: commands,
		pool:     pool,
	}

	go consumer.ConsumeForever()

	return consumer
}

func (r *RedisConsumer) ConsumeForever() {
	for {
		command := <-r.commands

		if DEBUG {
			log.Println("Dequeued one command in consumer")
		}

		conn, success := r.pool.Get()

		if success {
			command(conn)
			r.pool.Close(conn)
		} else {
			log.Println("Failed to get redis connection")
		}
	}
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"net/http"

	"github.com/julienschmidt/httprouter"
	"gopkg.in/mgo.v2"

	"github.com/herald-it/goncord/models"
	"github.com/herald-it/goncord/querying"
	. "github.com/herald-it/goncord/utils"
)

type ServiceController struct {
	session *mgo.Session
}

func (sc ServiceController) GetDB() *mgo.Database {
	return sc.session.DB("auth_service")
}

func NewServiceController(s *mgo.Session) *ServiceController {
	return &ServiceController{s}
}

func (sc ServiceController) IsValid(
	w http.ResponseWriter,
	r *http.Request,
	ps httprouter.Params) *HttpError {

	collect := sc.GetDB().C("token_dump")

	if err := r.ParseForm(); err != nil {
		return &HttpError{err, "Post form can not be parsed.", 500}
	}

	token := new(models.DumpToken)
	if err := Fill(token, r.PostForm); err != nil {
		return &HttpError{err, "Post form is not consistent with structure.", 500}
	}

	if token.Token == "" {
		return &HttpError{nil, "Invalid token value.", 500}
	}

	find_dump_token, err := querying.FindDumpToken(token, collect)
	if err != nil || find_dump_token == nil {
		return &HttpError{err, "Token not found.", 500}
	}

	usr := new(models.User)
	usr.Id = find_dump_token.UserId

	find_usr, err := querying.FindUser(usr, sc.GetDB().C("users"))
	if err != nil {
		return &HttpError{err, "User not found.", 500}
	}

	json_usr, err := json.Marshal(find_usr)
	if err != nil {
		return &HttpError{err, "User can not convert to json.", 500}
	}

	w.Write(json_usr)
	return nil
}
<commit_msg>Change constant table name on setting field.<commit_after>package controllers

import (
	"encoding/json"
	"net/http"

	"github.com/julienschmidt/httprouter"
	"gopkg.in/mgo.v2"

	"github.com/herald-it/goncord/models"
	"github.com/herald-it/goncord/querying"
	. "github.com/herald-it/goncord/utils"
)

type ServiceController struct {
	session *mgo.Session
}

func (sc ServiceController) GetDB() *mgo.Database {
	return sc.session.DB(models.Set.Database.DbName)
}

func NewServiceController(s *mgo.Session) *ServiceController {
	return &ServiceController{s}
}

func (sc ServiceController) IsValid(
	w http.ResponseWriter,
	r *http.Request,
	ps httprouter.Params) *HttpError {

	collect := sc.GetDB().C(models.Set.Database.TokenTable)

	if err := r.ParseForm(); err != nil {
		return &HttpError{err, "Post form can not be parsed.", 500}
	}

	token := new(models.DumpToken)
	if err := Fill(token, r.PostForm); err != nil {
		return &HttpError{err, "Post form is not consistent with structure.", 500}
	}

	if token.Token == "" {
		return &HttpError{nil, "Invalid token value.", 500}
	}

	find_dump_token, err := querying.FindDumpToken(token, collect)
	if err != nil || find_dump_token == nil {
		return &HttpError{err, "Token not found.", 500}
	}

	usr := new(models.User)
	usr.Id = find_dump_token.UserId

	find_usr, err := querying.FindUser(usr, sc.GetDB().C(models.Set.Database.UserTable))
	if err != nil {
		return &HttpError{err, "User not found.", 500}
	}

	json_usr, err := json.Marshal(find_usr)
	if err != nil {
		return &HttpError{err, "User can not convert to json.", 500}
	}

	w.Write(json_usr)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"path/filepath"

	"github.com/hwaf/hwaf/hlib"
)

// map of pkgname -> libname
//  if empty => ignore dep.
var g_pkg_map = map[string]string{
	"AtlasPolicy":        "",
	"AtlasCxxPolicy":     "",
	"AtlasFortranPolicy": "",
	"ExternalPolicy":     "",
	"GaudiInterface":     "GaudiKernel",
	"AtlasROOT":          "ROOT",
	"AtlasReflex":        "Reflex",
	"AtlasCLHEP":         "CLHEP",
	"AtlasPOOL":          "POOL",
	"AtlasCOOL":          "COOL",
	"AtlasCORAL":         "CORAL",
}

func find_tgt(wscript *hlib.Wscript_t, name string) (int, *hlib.Target_t) {
	wbld := &wscript.Build
	for i := range wbld.Targets {
		if wbld.Targets[i].Name == name {
			return i, &wbld.Targets[i]
		}
	}
	return -1, nil
}

func use_list(wscript *hlib.Wscript_t) []string {
	uses := []string{}
	for _, dep := range wscript.Package.Deps {
		pkg := filepath.Base(dep.Name)
		use_pkg, ok := g_pkg_map[pkg]
		if !ok {
			use_pkg = pkg
		}
		if use_pkg != "" {
			uses = append(uses, use_pkg)
		}
	}
	return uses
}

func cnv_atlas_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_component_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	libname := ""
	switch len(x.Args) {
	case 0:
		// component_library pattern
		libname = filepath.Base(wscript.Package.Name)
	default:
		// named_component_library pattern
		libname = x.Args[1]
	}
	itgt, tgt := find_tgt(wscript, libname)
	if itgt < 0 {
		wscript.Build.Targets = append(
			wscript.Build.Targets,
			hlib.Target_t{Name: libname},
		)
		itgt, tgt = find_tgt(wscript, libname)
	}
	tgt.Features = []string{"atlas_component"}
	tgt.Use = []hlib.Value{hlib.DefaultValue("uses", use_list(wscript))}

	fmt.Printf(">>> component [%v]...\n", *tgt)
	return nil
}

func cnv_atlas_dual_use_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_tpcnv_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_install_joboptions(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-jobos"}
	tgt.Features = []string{"atlas_install_joboptions"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"jobos",
		[]string{"share/*.py", "share/*.txt"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_python_modules(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-py"}
	tgt.Features = []string{"atlas_install_python_modules"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"python-files",
		[]string{"python/*.py"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_scripts(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-scripts"}
	tgt.Features = []string{"atlas_install_scripts"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"script-files",
		[]string{"scripts/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_xmls(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-xmls"}
	tgt.Features = []string{"atlas_install_xmls"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"xml-files",
		[]string{"xml/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_data(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-data"}
	tgt.Features = []string{"atlas_install_data"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"data-files",
		[]string{"data/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_java(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_dictionary(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> %v\n", x)
	return nil
}

// EOF
<commit_msg>cnv: index out of range<commit_after>package main

import (
	"fmt"
	"path/filepath"

	"github.com/hwaf/hwaf/hlib"
)

// map of pkgname -> libname
//  if empty => ignore dep.
var g_pkg_map = map[string]string{
	"AtlasPolicy":        "",
	"AtlasCxxPolicy":     "",
	"AtlasFortranPolicy": "",
	"ExternalPolicy":     "",
	"GaudiInterface":     "GaudiKernel",
	"AtlasROOT":          "ROOT",
	"AtlasReflex":        "Reflex",
	"AtlasCLHEP":         "CLHEP",
	"AtlasPOOL":          "POOL",
	"AtlasCOOL":          "COOL",
	"AtlasCORAL":         "CORAL",
}

func find_tgt(wscript *hlib.Wscript_t, name string) (int, *hlib.Target_t) {
	wbld := &wscript.Build
	for i := range wbld.Targets {
		if wbld.Targets[i].Name == name {
			return i, &wbld.Targets[i]
		}
	}
	return -1, nil
}

func use_list(wscript *hlib.Wscript_t) []string {
	uses := []string{}
	for _, dep := range wscript.Package.Deps {
		pkg := filepath.Base(dep.Name)
		use_pkg, ok := g_pkg_map[pkg]
		if !ok {
			use_pkg = pkg
		}
		if use_pkg != "" {
			uses = append(uses, use_pkg)
		}
	}
	return uses
}

func cnv_atlas_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_component_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	libname := ""
	switch len(x.Args) {
	case 0:
		// component_library pattern
		libname = filepath.Base(wscript.Package.Name)
	default:
		// named_component_library pattern
		libname = x.Args[0]
	}
	itgt, tgt := find_tgt(wscript, libname)
	if itgt < 0 {
		wscript.Build.Targets = append(
			wscript.Build.Targets,
			hlib.Target_t{Name: libname},
		)
		itgt, tgt = find_tgt(wscript, libname)
	}
	tgt.Features = []string{"atlas_component"}
	tgt.Use = []hlib.Value{hlib.DefaultValue("uses", use_list(wscript))}

	fmt.Printf(">>> component [%v]...\n", *tgt)
	return nil
}

func cnv_atlas_dual_use_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_tpcnv_library(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_install_joboptions(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-jobos"}
	tgt.Features = []string{"atlas_install_joboptions"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"jobos",
		[]string{"share/*.py", "share/*.txt"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_python_modules(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-py"}
	tgt.Features = []string{"atlas_install_python_modules"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"python-files",
		[]string{"python/*.py"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_scripts(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-scripts"}
	tgt.Features = []string{"atlas_install_scripts"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"script-files",
		[]string{"scripts/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_xmls(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-xmls"}
	tgt.Features = []string{"atlas_install_xmls"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"xml-files",
		[]string{"xml/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_data(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	pkgname := filepath.Base(wscript.Package.Name)
	tgt := hlib.Target_t{Name: pkgname + "-install-data"}
	tgt.Features = []string{"atlas_install_data"}
	tgt.Source = []hlib.Value{hlib.DefaultValue(
		"data-files",
		[]string{"data/*"},
	)}
	wscript.Build.Targets = append(wscript.Build.Targets, tgt)
	return nil
}

func cnv_atlas_install_java(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> [%s] \n", x.Name)
	return nil
}

func cnv_atlas_dictionary(wscript *hlib.Wscript_t, stmt Stmt) error {
	x := stmt.(*ApplyPattern)
	fmt.Printf(">>> %v\n", x)
	return nil
}

// EOF
<|endoftext|>


<commit_before>// TODO
// * syntax check mig.Action.Arguments before exec()
/* Mozilla InvestiGator Agent

Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Initial Developer of the Original Code is
Mozilla Corporation
Portions created by the Initial Developer are Copyright (C) 2013
the Initial Developer. All Rights Reserved.

Contributor(s):
Julien Vehent jvehent@mozilla.com [:ulfr]

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.
*/
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/streadway/amqp"
	"mig"
	"mig/modules/filechecker"
	"mig/pgp"
	"os"
	"os/exec"
	"strings"
	"time"
)

// build version
var version string

func main() {
	// parse command line argument
	// -m selects the mode {agent, filechecker, ...}
	var mode = flag.String("m", "agent", "module to run (eg. agent, filechecker)")
	flag.Parse()

	switch *mode {
	case "filechecker":
		// pass the rest of the arguments as a byte array
		// to the filechecker module
		var tmparg string
		for _, arg := range flag.Args() {
			tmparg = tmparg + arg
		}
		args := []byte(tmparg)
		fmt.Printf(filechecker.Run(args))
		os.Exit(0)
	case "agent":
		var ctx Context
		var err error

		// if init fails, sleep for one minute and try again. forever.
		for {
			ctx, err = Init()
			if err == nil {
				break
			}
			fmt.Println(err)
			fmt.Println("initialisation failed. sleep and retry.");
			time.Sleep(60 * time.Second)
		}

		// Goroutine that receives messages from AMQP
		go getCommands(ctx)

		// GoRoutine that parses and validates incoming commands
		go func(){
			for msg := range ctx.Channels.NewCommand {
				err = parseCommands(ctx, msg)
				if err != nil {
					log := mig.Log{Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
					err = ReportErrorToScheduler(log)
					if err != nil {
						ctx.Channels.Log <- mig.Log{Desc: "Unable to report failure to scheduler."}.Err()
					}
				}
			}
		}()

		// GoRoutine that executes commands that run as agent modules
		go func(){
			for cmd := range ctx.Channels.RunAgentCommand {
				err = runAgentModule(ctx, cmd)
				if err != nil {
					log := mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
				}
			}
		}()

		// GoRoutine that formats results and send them to scheduler
		go func() {
			for result := range ctx.Channels.Results {
				err = sendResults(ctx, result)
				if err != nil {
					// on failure, log and attempt to report it to the scheduler
					log := mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
					err = ReportErrorToScheduler(log)
					if err != nil {
						ctx.Channels.Log <- mig.Log{Desc: "Unable to report failure to scheduler."}.Err()
					}
				}
			}
		}()

		// GoRoutine that sends keepAlive messages to scheduler
		go keepAliveAgent(ctx)

		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Agent '%s' started.", ctx.Agent.QueueLoc)}

		// won't exit until this chan received something
		exitReason := <-ctx.Channels.Terminate
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Shutting down agent: '%v'", exitReason)}.Emerg()
		Destroy(ctx)
	}
}

// getCommands receives AMQP messages, and feed them to the action chan
func getCommands(ctx Context) (err error) {
	for m := range ctx.MQ.Bind.Chan {
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("received message '%s'", m.Body)}.Debug()

		// Ack this message only
		err := m.Ack(true)
		if err != nil {
			desc := fmt.Sprintf("Failed to acknowledge reception. Message will be ignored. Body: '%s'", m.Body)
			ctx.Channels.Log <- mig.Log{Desc: desc}.Err()
			continue
		}

		// pass it along
		ctx.Channels.NewCommand <- m.Body
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("received message. queued in position %d", len(ctx.Channels.NewCommand))}
	}
	return
}

// parseCommands transforms a message into a MIG Command struct, performs validation
// and run the command
func parseCommands(ctx Context, msg []byte) (err error) {
	var cmd mig.Command
	cmd.ID = 0	// safety net
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("parseCommands() -> %v", e)

			// if we have a command to return, update status and send back
			if cmd.ID > 0 {
				cmd.Results = mig.Log{Desc: fmt.Sprintf("%v", err)}.Err()
				cmd.Status = "failed"
				ctx.Channels.Results <- cmd
			}
		}
		ctx.Channels.Log <- mig.Log{Desc: "leaving parseCommands()"}.Debug()
	}()

	// unmarshal the received command into a command struct
	// if this fails, inform the scheduler and skip this message
	err = json.Unmarshal(msg, &cmd)
	if err != nil {
		panic(err)
	}

	// get an io.Reader from the public pgp key
	keyring, err := pgp.TransformArmoredPubKeyToKeyring(PUBLICPGPKEY)
	if err != nil {
		panic(err)
	}

	// Check the action syntax and signature
	err = cmd.Action.Validate(keyring)
	if err != nil {
		desc := fmt.Sprintf("action validation failed: %v", err)
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: desc}.Err()
		panic(desc)
	}

	// Expiration is verified by the Validate() call above, but we need
	// to verify the ScheduledDate ourselves
	if time.Now().Before(cmd.Action.ScheduledDate) {
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "action is scheduled for later"}.Err()
		panic("ScheduledDateInFuture")
	}

	switch cmd.Action.Order {
	case "filechecker":
		// send to the agent module execution path
		ctx.Channels.RunAgentCommand <- cmd
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "Command queued for execution"}
	case "shell":
		// send to the external command execution path
		ctx.Channels.RunExternalCommand <- cmd
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "Command queued for execution"}
	case "terminate":
		ctx.Channels.Terminate <- fmt.Errorf("Terminate order received from scheduler")
	default:
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("order '%s' is invalid", cmd.Action.Order)}
		panic("OrderNotUnderstood")
	}
	return
}

// runAgentModule is a generic command launcher for MIG modules that are
// built into the agent's binary. It handles commands timeout.
func runAgentModule(ctx Context, migCmd mig.Command) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("runCommand() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{Desc: "leaving runCommand()"}.Debug()
	}()

	// waiter is a channel that receives a message when the timeout expires
	waiter := make(chan error, 1)
	var out bytes.Buffer

	// Command arguments must be in json format
	tmpargs, err := json.Marshal(migCmd.Action.Arguments)
	if err != nil {
		panic(err)
	}

	// stringify the arguments
	cmdArgs := fmt.Sprintf("%s", tmpargs)

	// build the command line and execute
	cmd := exec.Command(os.Args[0], "-m", strings.ToLower(migCmd.Action.Order), cmdArgs)
	cmd.Stdout = &out
	if err := cmd.Start(); err != nil {
		panic(err)
	}

	// launch the waiter in a separate goroutine
	go func() {
		waiter <- cmd.Wait()
	}()

	select {

	// Timeout case: command has reached timeout, kill it
	case <-time.After(MODULETIMEOUT):
		ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command timed out. Killing it."}.Err()

		// update the command status and send the response back
		migCmd.Status = "timeout"
		ctx.Channels.Results <- migCmd

		// kill the command
		err := cmd.Process.Kill()
		if err != nil {
			panic(err)
		}
		<-waiter // allow goroutine to exit

	// Normal exit case: command has finished before the timeout
	case err := <-waiter:

		if err != nil {
			ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command failed."}.Err()
			// update the command status and send the response back
			migCmd.Status = "failed"
			ctx.Channels.Results <- migCmd
			panic(err)

		} else {
			ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command succeeded."}
			err = json.Unmarshal(out.Bytes(), &migCmd.Results)
			if err != nil {
				panic(err)
			}
			// mark command status as successfully completed
			migCmd.Status = "succeeded"
			// send the results back to the scheduler
			ctx.Channels.Results <- migCmd
		}
	}
	return
}

// sendResults builds a message body and send the command results back to the scheduler
func sendResults(ctx Context, result mig.Command) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("sendResults() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{Desc: "leaving sendResults()"}.Debug()
	}()

	result.AgentQueueLoc = ctx.Agent.QueueLoc
	body, err := json.Marshal(result)
	if err != nil {
		panic(err)
	}

	routingKey := fmt.Sprintf("mig.sched.%s", ctx.Agent.QueueLoc)
	err = publish(ctx, "mig", routingKey, body)
	if err != nil {
		panic(err)
	}

	return
}

// keepAliveAgent will send heartbeats messages to the scheduler at regular intervals
func keepAliveAgent(ctx Context) (err error) {
	// declare a keepalive message
	HeartBeat := mig.KeepAlive{
		Name:		ctx.Agent.Hostname,
		OS:		ctx.Agent.OS,
		Version:	version,
		QueueLoc:	ctx.Agent.QueueLoc,
		StartTime:	time.Now(),
	}

	// loop forever
	for {
		HeartBeat.HeartBeatTS = time.Now()
		body, err := json.Marshal(HeartBeat)
		if err != nil {
			desc := fmt.Sprintf("keepAliveAgent failed with error '%v'", err)
			ctx.Channels.Log <- mig.Log{Desc: desc}.Err()
		}
		desc := fmt.Sprintf("heartbeat '%s'", body)
		ctx.Channels.Log <- mig.Log{Desc: desc}.Debug()
		publish(ctx, "mig", "mig.keepalive", body)
		time.Sleep(ctx.Sleeper)
	}
	return
}

func ReportErrorToScheduler(log mig.Log) (err error){
	return
}

// publish is a generic function that sends messages to an AMQP exchange
func publish(ctx Context, exchange, routingKey string, body []byte) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("publish() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{Desc: "leaving publish()"}.Debug()
	}()

	msg := amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		ContentType:  "text/plain",
		Body:         []byte(body),
	}
	err = ctx.MQ.Chan.Publish(exchange, routingKey,
				true,  // is mandatory
				false, // is immediate
				msg)   // AMQP message
	if err != nil {
		panic(err)
	}
	desc := fmt.Sprintf("Message published to exchange '%s' with routing key '%s' and body '%s'", exchange, routingKey, msg.Body)
	ctx.Channels.Log <- mig.Log{Desc: desc}.Debug()
	return
}

<commit_msg>[minor] Agent: code cleanup<commit_after>// TODO
// * syntax check mig.Action.Arguments before exec()
/* Mozilla InvestiGator Agent

Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Initial Developer of the Original Code is
Mozilla Corporation
Portions created by the Initial Developer are Copyright (C) 2013
the Initial Developer. All Rights Reserved.

Contributor(s):
Julien Vehent jvehent@mozilla.com [:ulfr]

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.
*/
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/streadway/amqp"
	"mig"
	"mig/modules/filechecker"
	"mig/pgp"
	"os"
	"os/exec"
	"strings"
	"time"
)

// build version
var version string

func main() {
	// parse command line argument
	// -m selects the mode {agent, filechecker, ...}
	var mode = flag.String("m", "agent", "module to run (eg. agent, filechecker)")
	flag.Parse()

	switch *mode {

	case "filechecker":
		// pass the rest of the arguments as a byte array
		// to the filechecker module
		var tmparg string
		for _, arg := range flag.Args() {
			tmparg = tmparg + arg
		}
		args := []byte(tmparg)
		fmt.Printf(filechecker.Run(args))
		os.Exit(0)

	case "agent":
		var ctx Context
		var err error

		// if init fails, sleep for one minute and try again. forever.
		for {
			ctx, err = Init()
			if err == nil {
				break
			}
			fmt.Println(err)
			fmt.Println("initialisation failed. sleep and retry.");
			time.Sleep(60 * time.Second)
		}

		// Goroutine that receives messages from AMQP
		go getCommands(ctx)

		// GoRoutine that parses and validates incoming commands
		go func(){
			for msg := range ctx.Channels.NewCommand {
				err = parseCommands(ctx, msg)
				if err != nil {
					log := mig.Log{Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
				}
			}
		}()

		// GoRoutine that executes commands that run as agent modules
		go func(){
			for cmd := range ctx.Channels.RunAgentCommand {
				err = runAgentModule(ctx, cmd)
				if err != nil {
					log := mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
				}
			}
		}()

		// GoRoutine that formats results and send them to scheduler
		go func() {
			for result := range ctx.Channels.Results {
				err = sendResults(ctx, result)
				if err != nil {
					// on failure, log and attempt to report it to the scheduler
					log := mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
					ctx.Channels.Log <- log
				}
			}
		}()

		// GoRoutine that sends keepAlive messages to scheduler
		go keepAliveAgent(ctx)

		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Mozilla InvestiGator version %s: started agent %s", version, ctx.Agent.QueueLoc)}

		// won't exit until this chan received something
		exitReason := <-ctx.Channels.Terminate
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Shutting down agent: '%v'", exitReason)}.Emerg()
		Destroy(ctx)
	}
}

// getCommands receives AMQP messages, and feed them to the action chan
func getCommands(ctx Context) (err error) {
	for m := range ctx.MQ.Bind.Chan {
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("received message '%s'", m.Body)}.Debug()

		// Ack this message only
		err := m.Ack(true)
		if err != nil {
			desc := fmt.Sprintf("Failed to acknowledge reception. Message will be ignored. Body: '%s'", m.Body)
			ctx.Channels.Log <- mig.Log{Desc: desc}.Err()
			continue
		}

		// pass it along
		ctx.Channels.NewCommand <- m.Body
		ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("received message. queued in position %d", len(ctx.Channels.NewCommand))}
	}
	return
}

// parseCommands transforms a message into a MIG Command struct, performs validation
// and run the command
func parseCommands(ctx Context, msg []byte) (err error) {
	var cmd mig.Command
	cmd.ID = 0	// safety net
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("parseCommands() -> %v", e)

			// if we have a command to return, update status and send back
			if cmd.ID > 0 {
				cmd.Results = mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("%v", err)}.Err()
				cmd.Status = "failed"
				ctx.Channels.Results <- cmd
			}
		}
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "leaving parseCommands()"}.Debug()
	}()

	// unmarshal the received command into a command struct
	// if this fails, inform the scheduler and skip this message
	err = json.Unmarshal(msg, &cmd)
	if err != nil {
		panic(err)
	}

	// get an io.Reader from the public pgp key
	keyring, err := pgp.TransformArmoredPubKeyToKeyring(PUBLICPGPKEY)
	if err != nil {
		panic(err)
	}

	// Check the action syntax and signature
	err = cmd.Action.Validate(keyring)
	if err != nil {
		desc := fmt.Sprintf("action validation failed: %v", err)
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: desc}.Err()
		panic(desc)
	}

	// Expiration is verified by the Validate() call above, but we need
	// to verify the ScheduledDate ourselves
	if time.Now().Before(cmd.Action.ScheduledDate) {
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "action is scheduled for later"}.Err()
		panic("ScheduledDateInFuture")
	}

	switch cmd.Action.Order {
	case "filechecker":
		// send to the agent module execution path
		ctx.Channels.RunAgentCommand <- cmd
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "Command queued for execution"}
	case "shell":
		// send to the external command execution path
		ctx.Channels.RunExternalCommand <- cmd
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: "Command queued for execution"}
	case "terminate":
		ctx.Channels.Terminate <- fmt.Errorf("Terminate order received from scheduler")
	default:
		ctx.Channels.Log <- mig.Log{CommandID: cmd.ID, ActionID: cmd.Action.ID, Desc: fmt.Sprintf("order '%s' is invalid", cmd.Action.Order)}
		panic("OrderNotUnderstood")
	}
	return
}

// runAgentModule is a generic command launcher for MIG modules that are
// built into the agent's binary. It handles commands timeout.
func runAgentModule(ctx Context, migCmd mig.Command) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("runCommand() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{CommandID: migCmd.ID, ActionID: migCmd.Action.ID, Desc: "leaving runCommand()"}.Debug()
	}()

	ctx.Channels.Log <- mig.Log{CommandID: migCmd.ID, ActionID: migCmd.Action.ID, Desc: fmt.Sprintf("executing command '%s'", migCmd.Action.Order)}.Debug()
	// waiter is a channel that receives a message when the timeout expires
	waiter := make(chan error, 1)
	var out bytes.Buffer

	// Command arguments must be in json format
	tmpargs, err := json.Marshal(migCmd.Action.Arguments)
	if err != nil {
		panic(err)
	}

	// stringify the arguments
	cmdArgs := fmt.Sprintf("%s", tmpargs)

	// build the command line and execute
	cmd := exec.Command(os.Args[0], "-m", strings.ToLower(migCmd.Action.Order), cmdArgs)
	cmd.Stdout = &out
	if err := cmd.Start(); err != nil {
		panic(err)
	}

	// launch the waiter in a separate goroutine
	go func() {
		waiter <- cmd.Wait()
	}()

	select {

	// Timeout case: command has reached timeout, kill it
	case <-time.After(MODULETIMEOUT):
		ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command timed out. Killing it."}.Err()

		// update the command status and send the response back
		migCmd.Status = "timeout"
		ctx.Channels.Results <- migCmd

		// kill the command
		err := cmd.Process.Kill()
		if err != nil {
			panic(err)
		}
		<-waiter // allow goroutine to exit

	// Normal exit case: command has finished before the timeout
	case err := <-waiter:

		if err != nil {
			ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command failed."}.Err()
			// update the command status and send the response back
			migCmd.Status = "failed"
			ctx.Channels.Results <- migCmd
			panic(err)

		} else {
			ctx.Channels.Log <- mig.Log{ActionID: migCmd.Action.ID, CommandID: migCmd.ID, Desc: "command succeeded."}
			err = json.Unmarshal(out.Bytes(), &migCmd.Results)
			if err != nil {
				panic(err)
			}
			// mark command status as successfully completed
			migCmd.Status = "succeeded"
			// send the results back to the scheduler
			ctx.Channels.Results <- migCmd
		}
	}
	return
}

// sendResults builds a message body and send the command results back to the scheduler
func sendResults(ctx Context, result mig.Command) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("sendResults() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: "leaving sendResults()"}.Debug()
	}()

	ctx.Channels.Log <- mig.Log{CommandID: result.ID, ActionID: result.Action.ID, Desc: "sending command results"}
	result.AgentQueueLoc = ctx.Agent.QueueLoc
	body, err := json.Marshal(result)
	if err != nil {
		panic(err)
	}

	routingKey := fmt.Sprintf("mig.sched.%s", ctx.Agent.QueueLoc)
	err = publish(ctx, "mig", routingKey, body)
	if err != nil {
		panic(err)
	}

	return
}

// keepAliveAgent will send heartbeats messages to the scheduler at regular intervals
func keepAliveAgent(ctx Context) (err error) {
	// declare a keepalive message
	HeartBeat := mig.KeepAlive{
		Name:		ctx.Agent.Hostname,
		OS:		ctx.Agent.OS,
		Version:	version,
		QueueLoc:	ctx.Agent.QueueLoc,
		StartTime:	time.Now(),
	}

	// loop forever
	for {
		HeartBeat.HeartBeatTS = time.Now()
		body, err := json.Marshal(HeartBeat)
		if err != nil {
			desc := fmt.Sprintf("keepAliveAgent failed with error '%v'", err)
			ctx.Channels.Log <- mig.Log{Desc: desc}.Err()
		}
		desc := fmt.Sprintf("heartbeat '%s'", body)
		ctx.Channels.Log <- mig.Log{Desc: desc}.Debug()
		publish(ctx, "mig", "mig.keepalive", body)
		time.Sleep(ctx.Sleeper)
	}
	return
}

// publish is a generic function that sends messages to an AMQP exchange
func publish(ctx Context, exchange, routingKey string, body []byte) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("publish() -> %v", e)
		}
		ctx.Channels.Log <- mig.Log{Desc: "leaving publish()"}.Debug()
	}()

	msg := amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		ContentType:  "text/plain",
		Body:         []byte(body),
	}
	err = ctx.MQ.Chan.Publish(exchange, routingKey,
				true,  // is mandatory
				false, // is immediate
				msg)   // AMQP message
	if err != nil {
		panic(err)
	}
	desc := fmt.Sprintf("Message published to exchange '%s' with routing key '%s' and body '%s'", exchange, routingKey, msg.Body)
	ctx.Channels.Log <- mig.Log{Desc: desc}.Debug()
	return
}
<|endoftext|>


<commit_before>package monitor

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	_ "io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
)

type ipdata struct {
	Epoch       string `json:"Epoch"`
	SourceIP    string `json:"SourceIP"`
	SourcePort  string `json:"SourcePort"`
	ReceiveIP   string `json:"ReceiveIP"`
	ReceivePort string `json:"ReceivePort"`
}

/* The Monitor interface defines a series of methods which will be defined on
 * monitor structs. The Start method takes a channel to send messages over,
 * back to the configurator. The Stop method kills the process which is
 * performing the actual monitoring.
 */
type Monitor interface {
	Start(messages chan<- []byte, dockerComposeName string)
	Stop()
}

/* Store the MonitorName, which is the name of the program to execute, the
 * DockerDirs which are the directories to monitor which our Docker containers
 * of interest live in, and a pointer to the exec.Cmd struct which describes
 * the running command.
 */
type FSMonitor struct {
	MonitorName   string
	DockerDirs    []string
	fsWatcherProc *exec.Cmd
}

type ExecMonitor struct {
	MonitorName  string
	ContainerIds []string
}

type NetMonitor struct {
	MonitorName  string
	ContainerIds []string
}

// Memoize these. They're kind of expensive to get.
var dockerContainerIds = []string{}
var dockerContainerProcessIds = []string{}

func runCommandAndSlurpOutput(commandname string, args []string) ([]string, error) {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return nil, err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	output := []string{}
	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return nil, err
			}
		}
		if len(line) > 0 {
			output = append(output, string(line))
		}
	}
	return output, nil
}

func runCommandAndChannelOutput(commandname string, args []string, output chan<- []byte) error {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return err
			}
		}
		if len(line) > 0 {
			output <- line
		}
	}
	return nil
}

func getDockerContainerIds(dockerComposeName string) []string {
	if len(dockerContainerIds) != 0 {
		return dockerContainerIds
	}
	ids, err := runCommandAndSlurpOutput("docker-compose", []string{"ps", "-q"})
	if err != nil {
		panic(err)
	}
	dockerContainerIds = ids
	return ids
}

func getDockerContainerProcessIds(dockerComposeName string) []string {
	// Memoize this function
	if len(dockerContainerProcessIds) > 0 {
		return dockerContainerProcessIds
	}
	ids := getDockerContainerIds(dockerComposeName)
	arguments := []string{"inspect", "-f", "{{ .State.Pid }}"}
	arguments = append(arguments, ids...)
	output, err := runCommandAndSlurpOutput("docker", arguments)
	for i, out := range output {
		// FIXME: do we want to throw an error instead?
		if out == "0" {
			fmt.Println("Container isn't running: ", out)
			output = append(output[:i], output[i+1:]...)
			continue
		}
		output[i] = out
	}

	if err != nil {
		panic(err)
	}
	dockerContainerProcessIds = output
	return output
}

func (e ExecMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	ids := getDockerContainerIds(dockerComposeName)
	err := runCommandAndChannelOutput("cproc_monitor/proc", ids, messages)
	if err != nil {
		fmt.Println("Exec monitor failed")
		panic(err)
	}
}

func (n NetMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	output := getDockerContainerProcessIds(dockerComposeName)
	nmProcessorChan := make(chan []byte)
	go networkMonitorProcessor(messages, nmProcessorChan)
	for _, procId := range output {
		err := setSymlink(procId, procId)
		if err != nil {
			panic(err)
		}
		go startIPProcess(nmProcessorChan, procId, "tcpdump", "-tt", "-nn", "-i", "any", "-l")
	}
}

func networkMonitorProcessor(sending chan<- []byte, receiving <-chan []byte) error {
	platform := []byte{}
	carlson := []byte{}
	pattern := "(\\d+\\.\\d+) IP (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+) > (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+)"
	compiled_pattern, err := regexp.Compile(pattern)
	if err != nil {
		return err
	}
	for {
		platform = <-receiving
		matches := compiled_pattern.FindSubmatch(platform)
		if len(matches) == 0 {
			fmt.Println("Output did not match regex: ", string(platform))
			continue
		}
		carl := ipdata{
			Epoch:       string(matches[1]),
			SourceIP:    string(matches[2]),
			SourcePort:  string(matches[3]),
			ReceiveIP:   string(matches[4]),
			ReceivePort: string(matches[5]),
		}
		carlson, err = json.Marshal(carl)
		if err != nil {
			return err
		}
		fmt.Println("Sending json: ", string(carlson))
		sending <- carlson
	}

	return nil
}
func startIPProcess(messages chan<- []byte, procId string, watcherName string,
	watcherArgs ...string) {
	arguments := []string{"netns", "exec", procId, watcherName}
	arguments = append(arguments, watcherArgs...)
	err := runCommandAndChannelOutput("ip", arguments, messages)
	if err != nil {
		panic(err)
	}
}

func setSymlink(procId string, destination string) error {
	nameSpaceDir := "/var/run/netns/"
	target := nameSpaceDir + destination
	source := "/proc/" + procId + "/ns/net"

	err := os.MkdirAll(nameSpaceDir, 0777)
	if err != nil {
		return err
	}
	_, err = os.Stat(target)
	if err == nil {
		err = os.Remove(target)
		if err != nil {
			return err
		}
	}

	_, err = os.Stat(source)
	if os.IsNotExist(err) {
		return errors.New("Could not set symlink: Target directory " +
			target + " does not exist")
	}

	err = os.Symlink(source, target)
	return err
}

/* Returns the root directories of all of the Docker containers being monitored
 * as a list of strings.
 */
func (m FSMonitor) getDockerFSDirectory(dockerComposeName string) []string {
	pids := getDockerContainerProcessIds(dockerComposeName)
	roots := make([]string, len(pids))
	for i, pid := range pids {
		roots[i] = "/proc/" + pid + "/root"
	}
	return roots
}

/* Start the process running on the honeypot host to monitor the Docker
 * container. The Docker container's filesystem is mounted on the host. Find
 * the location of this filesystem with the getDockerFSDirectory function and
 * store it in the struct. Then create and start the process and forward
 * the output of the process on to the messages channel.
 */
func (m FSMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	m.DockerDirs = m.getDockerFSDirectory(dockerComposeName)
	// FIXME Make arguments configurable
	arguments := append([]string{"-ten"}, m.DockerDirs...)
	fmt.Println(m.MonitorName, arguments)
	m.fsWatcherProc = exec.Command(m.MonitorName, arguments...)
	defer m.fsWatcherProc.Wait()

	outpipe, err := m.fsWatcherProc.StdoutPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stdout pipe")
		panic(err)
	}
	stderr, err := m.fsWatcherProc.StderrPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stderr pipe")
		panic(err)
	}
	go io.Copy(os.Stderr, stderr)

	m.fsWatcherProc.Start()

	stdoutReader := bufio.NewReader(outpipe)

	for {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutReader.ReadLine()
			if err != nil {
				fmt.Println("File monitor stopped")
				panic(err)
			}
			fetch = f
			line = append(line, partial_line...)
		}
		messages <- line

	}
}

/* Stop the filesystem monitor. Kill the process monitoring the Docker
 * container's filesysem.
 */
func (m FSMonitor) Stop() {
	err := m.fsWatcherProc.Process.Kill()
	if err != nil {
		log.Fatal("Could not kill ", m.MonitorName, err)
	}
}
<commit_msg>Fix arguments in monitor<commit_after>package monitor

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	_ "io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
)

type ipdata struct {
	Epoch       string `json:"Epoch"`
	SourceIP    string `json:"SourceIP"`
	SourcePort  string `json:"SourcePort"`
	ReceiveIP   string `json:"ReceiveIP"`
	ReceivePort string `json:"ReceivePort"`
}

/* The Monitor interface defines a series of methods which will be defined on
 * monitor structs. The Start method takes a channel to send messages over,
 * back to the configurator. The Stop method kills the process which is
 * performing the actual monitoring.
 */
type Monitor interface {
	Start(messages chan<- []byte, dockerComposeName string)
	Stop()
}

/* Store the MonitorName, which is the name of the program to execute, the
 * DockerDirs which are the directories to monitor which our Docker containers
 * of interest live in, and a pointer to the exec.Cmd struct which describes
 * the running command.
 */
type FSMonitor struct {
	MonitorName   string
	DockerDirs    []string
	fsWatcherProc *exec.Cmd
}

type ExecMonitor struct {
	MonitorName  string
	ContainerIds []string
}

type NetMonitor struct {
	MonitorName  string
	ContainerIds []string
}

// Memoize these. They're kind of expensive to get.
var dockerContainerIds = []string{}
var dockerContainerProcessIds = []string{}

func runCommandAndSlurpOutput(commandname string, args []string) ([]string, error) {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return nil, err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	output := []string{}
	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return nil, err
			}
		}
		if len(line) > 0 {
			output = append(output, string(line))
		}
	}
	return output, nil
}

func runCommandAndChannelOutput(commandname string, args []string, output chan<- []byte) error {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return err
			}
		}
		if len(line) > 0 {
			output <- line
		}
	}
	return nil
}

func getDockerContainerIds(dockerComposeName string) []string {
	if len(dockerContainerIds) != 0 {
		return dockerContainerIds
	}
	ids, err := runCommandAndSlurpOutput("docker-compose", []string{"ps", "-q"})
	if err != nil {
		panic(err)
	}
	dockerContainerIds = ids
	return ids
}

func getDockerContainerProcessIds(dockerComposeName string) []string {
	// Memoize this function
	if len(dockerContainerProcessIds) > 0 {
		return dockerContainerProcessIds
	}
	ids := getDockerContainerIds(dockerComposeName)
	arguments := []string{"inspect", "-f", "{{ .State.Pid }}"}
	arguments = append(arguments, ids...)
	output, err := runCommandAndSlurpOutput("docker", arguments)
	for i, out := range output {
		// FIXME: do we want to throw an error instead?
		if out == "0" {
			fmt.Println("Container isn't running: ", out)
			output = append(output[:i], output[i+1:]...)
			continue
		}
		output[i] = out
	}

	if err != nil {
		panic(err)
	}
	dockerContainerProcessIds = output
	return output
}

func (e ExecMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	ids := getDockerContainerIds(dockerComposeName)
	err := runCommandAndChannelOutput("cproc_monitor/proc", ids, messages)
	if err != nil {
		fmt.Println("Exec monitor failed")
		panic(err)
	}
}

func (n NetMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	output := getDockerContainerProcessIds(dockerComposeName)
	nmProcessorChan := make(chan []byte)
	go networkMonitorProcessor(messages, nmProcessorChan)
	for _, procId := range output {
		err := setSymlink(procId, procId)
		if err != nil {
			panic(err)
		}
		go startIPProcess(nmProcessorChan, procId, "tcpdump", "-tt", "-nn", "-i", "any", "-l")
	}
}

func networkMonitorProcessor(sending chan<- []byte, receiving <-chan []byte) error {
	platform := []byte{}
	carlson := []byte{}
	pattern := "(\\d+\\.\\d+) IP (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+) > (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+)"
	compiled_pattern, err := regexp.Compile(pattern)
	if err != nil {
		return err
	}
	for {
		platform = <-receiving
		matches := compiled_pattern.FindSubmatch(platform)
		if len(matches) == 0 {
			fmt.Println("Output did not match regex: ", string(platform))
			continue
		}
		carl := ipdata{
			Epoch:       string(matches[1]),
			SourceIP:    string(matches[2]),
			SourcePort:  string(matches[3]),
			ReceiveIP:   string(matches[4]),
			ReceivePort: string(matches[5]),
		}
		carlson, err = json.Marshal(carl)
		if err != nil {
			return err
		}
		fmt.Println("Sending json: ", string(carlson))
		sending <- carlson
	}

	return nil
}
func startIPProcess(messages chan<- []byte, procId string, watcherName string,
	watcherArgs ...string) {
	arguments := []string{"netns", "exec", procId, watcherName}
	arguments = append(arguments, watcherArgs...)
	err := runCommandAndChannelOutput("ip", arguments, messages)
	if err != nil {
		panic(err)
	}
}

func setSymlink(procId string, destination string) error {
	nameSpaceDir := "/var/run/netns/"
	target := nameSpaceDir + destination
	source := "/proc/" + procId + "/ns/net"

	err := os.MkdirAll(nameSpaceDir, 0777)
	if err != nil {
		return err
	}
	_, err = os.Stat(target)
	if err == nil {
		err = os.Remove(target)
		if err != nil {
			return err
		}
	}

	_, err = os.Stat(source)
	if os.IsNotExist(err) {
		return errors.New("Could not set symlink: Target directory " +
			target + " does not exist")
	}

	err = os.Symlink(source, target)
	return err
}

/* Returns the root directories of all of the Docker containers being monitored
 * as a list of strings.
 */
func (m FSMonitor) getDockerFSDirectory(dockerComposeName string) []string {
	pids := getDockerContainerProcessIds(dockerComposeName)
	roots := make([]string, len(pids))
	for i, pid := range pids {
		roots[i] = "/proc/" + pid + "/root"
	}
	return roots
}

/* Start the process running on the honeypot host to monitor the Docker
 * container. The Docker container's filesystem is mounted on the host. Find
 * the location of this filesystem with the getDockerFSDirectory function and
 * store it in the struct. Then create and start the process and forward
 * the output of the process on to the messages channel.
 */
func (m FSMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	m.DockerDirs = m.getDockerFSDirectory(dockerComposeName)
	// FIXME Make arguments configurable
	arguments := append([]string{"-tea"}, m.DockerDirs...)
	fmt.Println(m.MonitorName, arguments)
	m.fsWatcherProc = exec.Command(m.MonitorName, arguments...)
	defer m.fsWatcherProc.Wait()

	outpipe, err := m.fsWatcherProc.StdoutPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stdout pipe")
		panic(err)
	}
	stderr, err := m.fsWatcherProc.StderrPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stderr pipe")
		panic(err)
	}
	go io.Copy(os.Stderr, stderr)

	m.fsWatcherProc.Start()

	stdoutReader := bufio.NewReader(outpipe)

	for {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutReader.ReadLine()
			if err != nil {
				fmt.Println("File monitor stopped")
				panic(err)
			}
			fetch = f
			line = append(line, partial_line...)
		}
		messages <- line

	}
}

/* Stop the filesystem monitor. Kill the process monitoring the Docker
 * container's filesysem.
 */
func (m FSMonitor) Stop() {
	err := m.fsWatcherProc.Process.Kill()
	if err != nil {
		log.Fatal("Could not kill ", m.MonitorName, err)
	}
}
<|endoftext|>