<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"

	"github.com/missdeer/golib/ebook"
	"github.com/missdeer/golib/fsutil"
	flags "github.com/jessevdk/go-flags"
)

// Options for all command line options, long name must match field name
type Options struct {
	ListenAndServe  string  `long:"listenAndServe" description:"set http listen and serve address, example: :8080"`
	Format          string  `short:"f" long:"format" description:"set generated file format, candidate values: mobi, epub, pdf"`
	List            bool    `short:"l" long:"list" description:"list supported novel websites"`
	LeftMargin      float64 `long:"leftMargin" description:"set left margin for PDF format"`
	TopMargin       float64 `long:"topMargin" description:"set top margin for PDF format"`
	PageWidth       float64 `long:"pageWidth" description:"set page width for PDF format(unit: mm)"`
	PageHeight      float64 `long:"pageHeight" description:"set page height for PDF format(unit: mm)"`
	PageType        string  `short:"p" long:"pageType" description:"set page type for PDF format, add suffix to output file name"`
	TitleFontSize   int     `long:"titleFontSize" description:"set title font point size for PDF format"`
	ContentFontSize int     `long:"contentFontSize" description:"set content font point size for PDF format"`
	LineSpacing     float64 `long:"lineSpacing" description:"set line spacing rate for PDF format"`
	PagesPerFile    int     `long:"pagesPerFile" description:"split the big single PDF file to several smaller PDF files, how many pages should be included in a file, 0 means don't split"`
	ChaptersPerFile int     `long:"chaptersPerFile" description:"split the big signle PDF file to several smaller PDF files, how many chapters should be included in a file, 0 means don't split"`
	FontFile        string  `long:"fontFile" description:"set TTF font file path"`
	RetryCount      int     `short:"r" long:"retryCount" description:"download retry count"`
	Timeout         int     `short:"t" long:"timeout" description:"download timeout seconds"`
	ParallelCount   int64   `long:"parallelCount" description:"parallel count for downloading"`
	ConfigFile      string  `short:"c" long:"configFile" description:"read configurations from local file"`
	OutputFile      string  `short:"o" long:"outputFile" description:"output file path"`
	FromChapter     int     `long:"fromChapter" description:"from chapter"`
	FromTitle       string  `long:"fromTitle" description:"from title"`
	ToChapter       int     `long:"toChapter" description:"to chapter"`
	ToTitle         string  `long:"toTitle" description:"to title"`
}

type tocPattern struct {
	host            string
	bookTitle       string
	bookTitlePos    int
	item            string
	articleTitlePos int
	articleURLPos   int
	isAbsoluteURL   bool
}

type pageContentMarker struct {
	host  string
	start []byte
	end   []byte
}

type novelSiteHandler struct {
	Title         string
	MatchPatterns []string
	Download      func(string, ebook.IBook)
}

var (
	novelSiteHandlers []*novelSiteHandler
	opts              Options
	sha1ver           string // sha1 revision used to build the program
	buildTime         string // when the executable was built
)

func registerNovelSiteHandler(h *novelSiteHandler) {
	novelSiteHandlers = append(novelSiteHandlers, h)
}

func listCommandHandler() {
	fmt.Println("内建支持小说网站：")
	for _, h := range novelSiteHandlers {
		urlMap := make(map[string]struct{})
		for _, p := range h.MatchPatterns {
			u := strings.Replace(p, `\`, ``, -1)
			idxStart := strings.Index(u, `www.`)
			idxEnd := strings.Index(u[idxStart:], `/`)
			u = u[:idxStart+idxEnd]
			urlMap[u] = struct{}{}
		}
		var urls []string
		for u := range urlMap {
			urls = append(urls, u)
		}
		fmt.Println("\t" + h.Title + ": " + strings.Join(urls, ", "))
	}
	fmt.Println("通过配置文件支持小说网站：")
	for _, h := range novelSiteConfigurations {
		fmt.Println("\t" + h.SiteName + ": " + h.Host)
	}
}

func parseConfigurations(content []byte, opts *Options) bool {
	var options map[string]interface{}
	if err := json.Unmarshal(content, &options); err != nil {
		log.Println("unmarshall configurations failed", err)
		return false
	}

	oe := reflect.ValueOf(opts).Elem()
	for i := 0; i < oe.NumField(); i++ {
		fieldName := oe.Type().Field(i).Name
		key := strings.ToLower(fieldName[:1]) + fieldName[1:]
		if f, ok := options[key]; ok {
			of := oe.Field(i)
			switch of.Kind() {
			case reflect.String:
				if v := f.(string); len(v) > 0 {
					of.SetString(v)
				}
			case reflect.Float64:
				if v := f.(float64); v > 0 {
					of.SetFloat(v)
				}
			case reflect.Int, reflect.Int64:
				if v := f.(float64); v > 0 {
					of.SetInt(int64(v))
				}
			}
		}
	}
	return true
}

func readRemotePreset(opts *Options) bool {
	u := "https://cdn.jsdelivr.net/gh/missdeer/getnovel/preset/" + opts.ConfigFile
	client := &http.Client{}
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("Could not parse preset request:", err)
		return false
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Println("Could not send request:", err)
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("response not 200:", resp.StatusCode, resp.Status)
		return false
	}

	c, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("reading content failed")
		return false
	}

	return parseConfigurations(c, opts)
}

func readLocalConfigFile(opts *Options) bool {
	configFile := opts.ConfigFile
	if b, e := fsutil.FileExists(configFile); e != nil || !b {
		configFile = filepath.Join("preset", opts.ConfigFile)
		if b, e = fsutil.FileExists(configFile); e != nil || !b {
			log.Println("cannot find configuration file", opts.ConfigFile, "on local file system")
			return false
		}
	}

	contentFd, err := os.OpenFile(configFile, os.O_RDONLY, 0644)
	if err != nil {
		log.Println("opening config file", configFile, "for reading failed", err)
		return false
	}

	contentC, err := ioutil.ReadAll(contentFd)
	contentFd.Close()
	if err != nil {
		log.Println("reading config file", configFile, "failed", err)
		return false
	}

	return parseConfigurations(contentC, opts)
}

func downloadBook(novelURL string, ch chan bool) {
	for _, h := range novelSiteHandlers {
		for _, pattern := range h.MatchPatterns {
			r, _ := regexp.Compile(pattern)
			if r.MatchString(novelURL) {
				gen := ebook.NewBook(opts.Format)
				gen.SetFontSize(opts.TitleFontSize, opts.ContentFontSize)
				gen.SetLineSpacing(opts.LineSpacing)
				gen.PagesPerFile(opts.PagesPerFile)
				gen.ChaptersPerFile(opts.ChaptersPerFile)
				gen.SetMargins(opts.LeftMargin, opts.TopMargin)
				gen.SetPageType(opts.PageType)
				gen.SetPageSize(opts.PageWidth, opts.PageHeight)
				gen.SetFontFile(opts.FontFile)
				gen.Output(opts.OutputFile)
				gen.Info()
				h.Download(novelURL, gen)
				fmt.Println("downloaded", novelURL)
				ch <- true
				return
			}
		}
	}

	u, _ := url.Parse(novelURL)
	for _, h := range novelSiteConfigurations {
		theURL, _ := url.Parse(h.Host)
		if u.Scheme == theURL.Scheme && u.Host == theURL.Host {
			gen := ebook.NewBook(opts.Format)
			gen.SetFontSize(opts.TitleFontSize, opts.ContentFontSize)
			gen.SetLineSpacing(opts.LineSpacing)
			gen.PagesPerFile(opts.PagesPerFile)
			gen.ChaptersPerFile(opts.ChaptersPerFile)
			gen.SetMargins(opts.LeftMargin, opts.TopMargin)
			gen.SetPageType(opts.PageType)
			gen.SetPageSize(opts.PageWidth, opts.PageHeight)
			gen.SetFontFile(opts.FontFile)
			gen.Output(opts.OutputFile)
			gen.Info()
			h.Download(novelURL, gen)
			fmt.Println("downloaded", novelURL)
			ch <- true
			return
		}
	}
	fmt.Println("not downloaded", novelURL)
	ch <- false
}

func main() {
	fmt.Println("getnovel SHA1:", sha1ver, "build at", buildTime)
	if len(os.Args) < 2 {
		fmt.Println("使用方法：\n\tgetnovel 小说目录网址")
		listCommandHandler()
		return
	}

	opts = Options{
		Format:          "mobi",
		List:            false,
		LeftMargin:      10,
		TopMargin:       10,
		PageHeight:      841.89,
		PageWidth:       595.28,
		TitleFontSize:   24,
		ContentFontSize: 18,
		LineSpacing:     1.2,
		PagesPerFile:    0,
		ChaptersPerFile: 0,
		FontFile:        filepath.Join("fonts", "CustomFont.ttf"),
		RetryCount:      3,
		Timeout:         60,
		ParallelCount:   10,
	}

	args, err := flags.Parse(&opts)
	if err != nil {
		log.Fatalln("parsing flags failed", err)
		return
	}

	readNovelSiteConfigurations()

	if opts.List {
		listCommandHandler()
		return
	}

	if opts.ListenAndServe != "" {
		dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
		if err != nil {
			log.Fatal(err)
		}
		ifaces, err := net.Interfaces()
		var ips []string
		for _, i := range ifaces {

			addrs, err := i.Addrs()
			if err != nil {
				fmt.Println(err)
				continue
			}
			for _, addr := range addrs {
				switch v := addr.(type) {
				case *net.IPNet:
					if v.IP.IsLoopback() || v.IP.IsLinkLocalMulticast() || v.IP.IsLinkLocalUnicast() {
						break
					}
					ips = append(ips, "\t"+v.IP.String())
				case *net.IPAddr:
					if v.IP.IsLoopback() || v.IP.IsLinkLocalMulticast() || v.IP.IsLinkLocalUnicast() {
						break
					}
					ips = append(ips, "\t"+v.IP.String())
				}
			}
		}
		fmt.Println("Local IP:")
		fmt.Println(strings.Join(ips, "\n"))
		fmt.Println("starting http server on", opts.ListenAndServe)
		log.Fatal(http.ListenAndServe(opts.ListenAndServe, http.FileServer(http.Dir(dir))))
		return
	}

	if opts.ConfigFile != "" {
		if !readLocalConfigFile(&opts) && !readRemotePreset(&opts) {
			return
		}
	}

	downloadedChannel := make(chan bool)
	donwloadCount := 0
	for _, novelURL := range args {
		_, e := url.Parse(novelURL)
		if e != nil {
			fmt.Println("invalid URL", novelURL)
			continue
		}
		donwloadCount++
		go downloadBook(novelURL, downloadedChannel)
	}

	downloaded := false
	for i := 0; i < donwloadCount; i++ {
		ch := <-downloadedChannel
		downloaded = (downloaded || ch)
	}

	if !downloaded {
		fmt.Println("使用方法：\n\tgetnovel 小说目录网址")
		listCommandHandler()
	}
}
<commit_msg>(*)fix typo<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"

	"github.com/jessevdk/go-flags"
	"github.com/missdeer/golib/ebook"
	"github.com/missdeer/golib/fsutil"
)

// Options for all command line options, long name must match field name
type Options struct {
	ListenAndServe  string  `long:"listenAndServe" description:"set http listen and serve address, example: :8080"`
	Format          string  `short:"f" long:"format" description:"set generated file format, candidate values: mobi, epub, pdf"`
	List            bool    `short:"l" long:"list" description:"list supported novel websites"`
	LeftMargin      float64 `long:"leftMargin" description:"set left margin for PDF format"`
	TopMargin       float64 `long:"topMargin" description:"set top margin for PDF format"`
	PageWidth       float64 `long:"pageWidth" description:"set page width for PDF format(unit: mm)"`
	PageHeight      float64 `long:"pageHeight" description:"set page height for PDF format(unit: mm)"`
	PageType        string  `short:"p" long:"pageType" description:"set page type for PDF format, add suffix to output file name"`
	TitleFontSize   int     `long:"titleFontSize" description:"set title font point size for PDF format"`
	ContentFontSize int     `long:"contentFontSize" description:"set content font point size for PDF format"`
	LineSpacing     float64 `long:"lineSpacing" description:"set line spacing rate for PDF format"`
	PagesPerFile    int     `long:"pagesPerFile" description:"split the big single PDF file to several smaller PDF files, how many pages should be included in a file, 0 means don't split"`
	ChaptersPerFile int     `long:"chaptersPerFile" description:"split the big single PDF file to several smaller PDF files, how many chapters should be included in a file, 0 means don't split"`
	FontFile        string  `long:"fontFile" description:"set TTF font file path"`
	RetryCount      int     `short:"r" long:"retryCount" description:"download retry count"`
	Timeout         int     `short:"t" long:"timeout" description:"download timeout seconds"`
	ParallelCount   int64   `long:"parallelCount" description:"parallel count for downloading"`
	ConfigFile      string  `short:"c" long:"configFile" description:"read configurations from local file"`
	OutputFile      string  `short:"o" long:"outputFile" description:"output file path"`
	FromChapter     int     `long:"fromChapter" description:"from chapter"`
	FromTitle       string  `long:"fromTitle" description:"from title"`
	ToChapter       int     `long:"toChapter" description:"to chapter"`
	ToTitle         string  `long:"toTitle" description:"to title"`
}

type tocPattern struct {
	host            string
	bookTitle       string
	bookTitlePos    int
	item            string
	articleTitlePos int
	articleURLPos   int
	isAbsoluteURL   bool
}

type pageContentMarker struct {
	host  string
	start []byte
	end   []byte
}

type novelSiteHandler struct {
	Title         string
	MatchPatterns []string
	Download      func(string, ebook.IBook)
}

var (
	novelSiteHandlers []*novelSiteHandler
	opts              Options
	sha1ver           string // sha1 revision used to build the program
	buildTime         string // when the executable was built
)

func registerNovelSiteHandler(h *novelSiteHandler) {
	novelSiteHandlers = append(novelSiteHandlers, h)
}

func listCommandHandler() {
	fmt.Println("内建支持小说网站：")
	for _, h := range novelSiteHandlers {
		urlMap := make(map[string]struct{})
		for _, p := range h.MatchPatterns {
			u := strings.Replace(p, `\`, ``, -1)
			idxStart := strings.Index(u, `www.`)
			idxEnd := strings.Index(u[idxStart:], `/`)
			u = u[:idxStart+idxEnd]
			urlMap[u] = struct{}{}
		}
		var urls []string
		for u := range urlMap {
			urls = append(urls, u)
		}
		fmt.Println("\t" + h.Title + ": " + strings.Join(urls, ", "))
	}
	fmt.Println("通过配置文件支持小说网站：")
	for _, h := range novelSiteConfigurations {
		fmt.Println("\t" + h.SiteName + ": " + h.Host)
	}
}

func parseConfigurations(content []byte, opts *Options) bool {
	var options map[string]interface{}
	if err := json.Unmarshal(content, &options); err != nil {
		log.Println("unmarshal configurations failed", err)
		return false
	}

	oe := reflect.ValueOf(opts).Elem()
	for i := 0; i < oe.NumField(); i++ {
		fieldName := oe.Type().Field(i).Name
		key := strings.ToLower(fieldName[:1]) + fieldName[1:]
		if f, ok := options[key]; ok {
			of := oe.Field(i)
			switch of.Kind() {
			case reflect.String:
				if v := f.(string); len(v) > 0 {
					of.SetString(v)
				}
			case reflect.Float64:
				if v := f.(float64); v > 0 {
					of.SetFloat(v)
				}
			case reflect.Int, reflect.Int64:
				if v := f.(float64); v > 0 {
					of.SetInt(int64(v))
				}
			}
		}
	}
	return true
}

func readRemotePreset(opts *Options) bool {
	u := "https://cdn.jsdelivr.net/gh/missdeer/getnovel/preset/" + opts.ConfigFile
	client := &http.Client{}
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("Could not parse preset request:", err)
		return false
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Println("Could not send request:", err)
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("response not 200:", resp.StatusCode, resp.Status)
		return false
	}

	c, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("reading content failed")
		return false
	}

	return parseConfigurations(c, opts)
}

func readLocalConfigFile(opts *Options) bool {
	configFile := opts.ConfigFile
	if b, e := fsutil.FileExists(configFile); e != nil || !b {
		configFile = filepath.Join("preset", opts.ConfigFile)
		if b, e = fsutil.FileExists(configFile); e != nil || !b {
			log.Println("cannot find configuration file", opts.ConfigFile, "on local file system")
			return false
		}
	}

	contentFd, err := os.OpenFile(configFile, os.O_RDONLY, 0644)
	if err != nil {
		log.Println("opening config file", configFile, "for reading failed", err)
		return false
	}

	contentC, err := ioutil.ReadAll(contentFd)
	contentFd.Close()
	if err != nil {
		log.Println("reading config file", configFile, "failed", err)
		return false
	}

	return parseConfigurations(contentC, opts)
}

func downloadBook(novelURL string, ch chan bool) {
	for _, h := range novelSiteHandlers {
		for _, pattern := range h.MatchPatterns {
			r, _ := regexp.Compile(pattern)
			if r.MatchString(novelURL) {
				gen := ebook.NewBook(opts.Format)
				gen.SetFontSize(opts.TitleFontSize, opts.ContentFontSize)
				gen.SetLineSpacing(opts.LineSpacing)
				gen.PagesPerFile(opts.PagesPerFile)
				gen.ChaptersPerFile(opts.ChaptersPerFile)
				gen.SetMargins(opts.LeftMargin, opts.TopMargin)
				gen.SetPageType(opts.PageType)
				gen.SetPageSize(opts.PageWidth, opts.PageHeight)
				gen.SetFontFile(opts.FontFile)
				gen.Output(opts.OutputFile)
				gen.Info()
				h.Download(novelURL, gen)
				fmt.Println("downloaded", novelURL)
				ch <- true
				return
			}
		}
	}

	u, _ := url.Parse(novelURL)
	for _, h := range novelSiteConfigurations {
		theURL, _ := url.Parse(h.Host)
		if u.Scheme == theURL.Scheme && u.Host == theURL.Host {
			gen := ebook.NewBook(opts.Format)
			gen.SetFontSize(opts.TitleFontSize, opts.ContentFontSize)
			gen.SetLineSpacing(opts.LineSpacing)
			gen.PagesPerFile(opts.PagesPerFile)
			gen.ChaptersPerFile(opts.ChaptersPerFile)
			gen.SetMargins(opts.LeftMargin, opts.TopMargin)
			gen.SetPageType(opts.PageType)
			gen.SetPageSize(opts.PageWidth, opts.PageHeight)
			gen.SetFontFile(opts.FontFile)
			gen.Output(opts.OutputFile)
			gen.Info()
			h.Download(novelURL, gen)
			fmt.Println("downloaded", novelURL)
			ch <- true
			return
		}
	}
	fmt.Println("not downloaded", novelURL)
	ch <- false
}

func main() {
	fmt.Println("getnovel SHA1:", sha1ver, "build at", buildTime)
	if len(os.Args) < 2 {
		fmt.Println("使用方法：\n\tgetnovel 小说目录网址")
		listCommandHandler()
		return
	}

	opts = Options{
		Format:          "mobi",
		List:            false,
		LeftMargin:      10,
		TopMargin:       10,
		PageHeight:      841.89,
		PageWidth:       595.28,
		TitleFontSize:   24,
		ContentFontSize: 18,
		LineSpacing:     1.2,
		PagesPerFile:    0,
		ChaptersPerFile: 0,
		FontFile:        filepath.Join("fonts", "CustomFont.ttf"),
		RetryCount:      3,
		Timeout:         60,
		ParallelCount:   10,
	}

	args, err := flags.Parse(&opts)
	if err != nil {
		log.Fatalln("parsing flags failed", err)
		return
	}

	readNovelSiteConfigurations()

	if opts.List {
		listCommandHandler()
		return
	}

	if opts.ListenAndServe != "" {
		dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
		if err != nil {
			log.Fatal(err)
		}
		ifaces, err := net.Interfaces()
		var ips []string
		for _, i := range ifaces {

			addrs, err := i.Addrs()
			if err != nil {
				fmt.Println(err)
				continue
			}
			for _, addr := range addrs {
				switch v := addr.(type) {
				case *net.IPNet:
					if v.IP.IsLoopback() || v.IP.IsLinkLocalMulticast() || v.IP.IsLinkLocalUnicast() {
						break
					}
					ips = append(ips, "\t"+v.IP.String())
				case *net.IPAddr:
					if v.IP.IsLoopback() || v.IP.IsLinkLocalMulticast() || v.IP.IsLinkLocalUnicast() {
						break
					}
					ips = append(ips, "\t"+v.IP.String())
				}
			}
		}
		fmt.Println("Local IP:")
		fmt.Println(strings.Join(ips, "\n"))
		fmt.Println("starting http server on", opts.ListenAndServe)
		log.Fatal(http.ListenAndServe(opts.ListenAndServe, http.FileServer(http.Dir(dir))))
		return
	}

	if opts.ConfigFile != "" {
		if !readLocalConfigFile(&opts) && !readRemotePreset(&opts) {
			return
		}
	}

	downloadedChannel := make(chan bool)
	downloadCount := 0
	for _, novelURL := range args {
		_, e := url.Parse(novelURL)
		if e != nil {
			fmt.Println("invalid URL", novelURL)
			continue
		}
		downloadCount++
		go downloadBook(novelURL, downloadedChannel)
	}

	downloaded := false
	for i := 0; i < downloadCount; i++ {
		ch := <-downloadedChannel
		downloaded = downloaded || ch
	}

	if !downloaded {
		fmt.Println("使用方法：\n\tgetnovel 小说目录网址")
		listCommandHandler()
	}
}
<|endoftext|>


<commit_before>package main

import (
  "flag"
  "fmt"
  "github.com/jzelinskie/geddit"
  "gopkg.in/yaml.v2"
  "io/ioutil"
  "log"
  "os"
  "os/signal"
  "syscall"
  "time"
)

type Bot struct {
  Username string
  Password string
}

func main() {

  // Using gist 'aggrolite/160d5be23adb9e597553', 'go-reddit-bot' as resources
  var bot Bot
  var config string
  var outfile string

  flag.StringVar(&config, "config", "bot.yml", "Path to YAML config file.")
  flag.StringVar(&outfile, "output", "links.txt", "Path to an output text file.")
  flag.Parse()

  // Read config file.
  config_data, err := ioutil.ReadFile(config)
  if err != nil {
    log.Fatal(err)
  }
  err = yaml.Unmarshal(config_data, &bot)
  if err != nil {
    log.Fatal(err)
  }

  // Use own reddit API wrapper once done.
  session, err := geddit.NewLoginSession(
    bot.Username,
    bot.Password,
    "Link Saver written in Go by u/Alexius-CA",
  )
  if err != nil {
    log.Fatal(err)
  }

  // Query options if needed (blank right now).
  //opts := geddit.ListingOptions{}

  ticker := time.NewTicker(30 * time.Second)
  quit := make(chan os.Signal)
  signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

  // Acquire saved links and write to file every 1 hour.
  go func() {
    for {
      select {
      case <-ticker.C:
        saved, err := session.MySaved(geddit.NewSubmissions, "")
        write := ""
        for _, s := range saved {
          write += fmt.Sprintf("#%s\n%s\n", s.Subreddit, s.Permalink)
        }
        err = ioutil.WriteFile(outfile, []byte(write), 0644)
        if err != nil {
          log.Print(err)
        }
      case <-quit:
        ticker.Stop()
        return
      }
    }
  }()
  <-quit

}
<commit_msg>Baka.<commit_after>package main

import (
  "flag"
  "fmt"
  "github.com/jzelinskie/geddit"
  "gopkg.in/yaml.v2"
  "io/ioutil"
  "log"
  "os"
  "os/signal"
  "syscall"
  "time"
)

type Bot struct {
  Username string
  Password string
}

func main() {

  // Using gist 'aggrolite/160d5be23adb9e597553', 'go-reddit-bot' as resources
  var bot Bot
  var config string
  var outfile string

  flag.StringVar(&config, "config", "bot.yml", "Path to YAML config file.")
  flag.StringVar(&outfile, "output", "links.txt", "Path to an output text file.")
  flag.Parse()

  // Read config file.
  config_data, err := ioutil.ReadFile(config)
  if err != nil {
    log.Fatal(err)
  }
  err = yaml.Unmarshal(config_data, &bot)
  if err != nil {
    log.Fatal(err)
  }

  // Use own reddit API wrapper once done.
  session, err := geddit.NewLoginSession(
    bot.Username,
    bot.Password,
    "Saved Links Bot (golang) by u/Alexius-CA @ github.com:AlexSafatli/go-reddit-links-bot",
  )
  if err != nil {
    log.Fatal(err)
  }

  // Query options if needed (blank right now).
  //opts := geddit.ListingOptions{}

  ticker := time.NewTicker(30 * time.Second)
  quit := make(chan os.Signal)
  signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

  // Acquire saved links and write to file every 1 hour.
  go func() {
    for {
      select {
      case <-ticker.C:
        saved, err := session.MySaved(geddit.NewSubmissions, "")
        write := ""
        for _, s := range saved {
          write += fmt.Sprintf("#%s\n%s\n", s.Subreddit, s.Permalink)
        }
        err = ioutil.WriteFile(outfile, []byte(write), 0644)
        if err != nil {
          log.Print(err)
        }
      case <-quit:
        ticker.Stop()
        return
      }
    }
  }()
  <-quit

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	flags "github.com/jessevdk/go-flags"
	"github.com/cactus/gologit"
	"os"
	"runtime"
	"strconv"
	"syscall"
)

const VERSION = "0.0.1"

func chanByteReader(msg *SyslogMsg) {
	fmt.Printf("<%d> %s\n", msg.Priority, msg.Msg)
}

func main() {
	var gmx int
	if gmxEnv := os.Getenv("GOMAXPROCS"); gmxEnv != "" {
		gmx, _ = strconv.Atoi(gmxEnv)
	} else {
		gmx = runtime.NumCPU()
	}
	runtime.GOMAXPROCS(gmx)

	// command line flags
	var opts struct {
		BindTCP string `short:"t" long:"listen-tcp" default:"" description:"TCP address:port to listen to"`
		BindUDP string `short:"u" long:"listen-udp" default:"" description:"UDP address:port to listen to"`
		Verbose bool   `short:"v" long:"verbose" description:"Show verbose (debug) log level output"`
		Version bool   `short:"V" long:"version" description:"print version and exit"`
	}

	// parse said flags
	_, err := flags.Parse(&opts)
	if err != nil {
		if e, ok := err.(*flags.Error); ok {
			if e.Type == flags.ErrHelp {
				os.Exit(0)
			}
		}
		os.Exit(1)
	}

	if opts.Version {
		fmt.Printf("slurpy-%s (%s,%s-%s)\n", VERSION, runtime.Version(), runtime.Compiler, runtime.GOARCH)
		os.Exit(0)
	}

	if opts.BindTCP == "" && opts.BindUDP == "" {
		fmt.Println("No listen ports defined. Exiting.")
		os.Exit(1)
	}

	// set logger debug level and start toggle on signal handler
	logger := gologit.Logger
	logger.Set(opts.Verbose)
	logger.Debugln("Debug logging enabled")
	logger.ToggleOnSignal(syscall.SIGUSR1)

	if opts.BindTCP != "" {
		tcpsrv, err := ListenTCP("tcp", opts.BindTCP, chanByteReader)
		if err != nil {
			gologit.Fatal(err)
		}
		defer tcpsrv.Close()
	}

	if opts.BindUDP != "" {
		gologit.Println("Starting udp server on", opts.BindUDP)
		udpsrv, err := ListenUDP("udp", opts.BindUDP, chanByteReader)
		if err != nil {
			gologit.Fatal(err)
		}
		defer udpsrv.Close()
	}
	select {}
}
<commit_msg>add "nice" shutdown<commit_after>package main

import (
	"fmt"
	flags "github.com/jessevdk/go-flags"
	"github.com/cactus/gologit"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"syscall"
)

const VERSION = "0.0.1"

func chanByteReader(msg *SyslogMsg) {
	fmt.Printf("<%d> %s\n", msg.Priority, msg.Msg)
}

func main() {
	var gmx int
	if gmxEnv := os.Getenv("GOMAXPROCS"); gmxEnv != "" {
		gmx, _ = strconv.Atoi(gmxEnv)
	} else {
		gmx = runtime.NumCPU()
	}
	runtime.GOMAXPROCS(gmx)

	// command line flags
	var opts struct {
		BindTCP string `short:"t" long:"listen-tcp" default:"" description:"TCP address:port to listen to"`
		BindUDP string `short:"u" long:"listen-udp" default:"" description:"UDP address:port to listen to"`
		Verbose bool   `short:"v" long:"verbose" description:"Show verbose (debug) log level output"`
		Version bool   `short:"V" long:"version" description:"print version and exit"`
	}

	// parse said flags
	_, err := flags.Parse(&opts)
	if err != nil {
		if e, ok := err.(*flags.Error); ok {
			if e.Type == flags.ErrHelp {
				os.Exit(0)
			}
		}
		os.Exit(1)
	}

	if opts.Version {
		fmt.Printf("slurpy-%s (%s,%s-%s)\n", VERSION, runtime.Version(), runtime.Compiler, runtime.GOARCH)
		os.Exit(0)
	}

	if opts.BindTCP == "" && opts.BindUDP == "" {
		fmt.Println("No listen ports defined. Exiting.")
		os.Exit(1)
	}

	// set logger debug level and start toggle on signal handler
	logger := gologit.Logger
	logger.Set(opts.Verbose)
	logger.Debugln("Debug logging enabled")
	logger.ToggleOnSignal(syscall.SIGUSR1)

	if opts.BindTCP != "" {
		gologit.Println("Starting tcp server on", opts.BindTCP)
		tcpsrv, err := ListenTCP("tcp", opts.BindTCP, chanByteReader)
		if err != nil {
			gologit.Fatal(err)
		}
		defer tcpsrv.Close()
	}

	if opts.BindUDP != "" {
		gologit.Println("Starting udp server on", opts.BindUDP)
		udpsrv, err := ListenUDP("udp", opts.BindUDP, chanByteReader)
		if err != nil {
			gologit.Fatal(err)
		}
		defer udpsrv.Close()
	}
	closeSig := make(chan os.Signal, 1)
	signal.Notify(closeSig, os.Interrupt)
	s := <-closeSig
	gologit.Printf("got signal: %s. Closing.\n", s)
	return
}
<|endoftext|>


<commit_before>/* ELF tree - Tree viewer for ELF library dependency */
package main

import (
	"bufio"
	"debug/elf"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type DepsNode struct {
	name   string
	parent *DepsNode
	child  []*DepsNode
	depth  int
}

type DynInfo struct {
	tag elf.DynTag
	val interface{}
}

type DepsInfo struct {
	path   string
	mach   elf.Machine
	bits   elf.Class
	endian binary.ByteOrder
	kind   elf.Type
	abi    elf.OSABI
	ver    uint8

	libs []string
	dsym []elf.ImportedSymbol
	syms []elf.Symbol
	prog []*elf.Prog
	dyns []DynInfo
}

var (
	deps      map[string]DepsInfo
	deps_list []*DepsNode
	deps_root *DepsNode
	deflib    []string
	envlib    string
	conflib   []string
)

// command-line options
var (
	verbose  bool
	showPath bool
	showTui  bool
)

func readLdSoConf(name string, libpath []string) {
	f, err := os.Open(name)
	if err != nil {
		return
	}
	defer f.Close()

	s := bufio.NewScanner(f)
	for s.Scan() {
		t := s.Text()

		if len(strings.TrimSpace(t)) == 0 {
			continue
		}
		if strings.HasPrefix(t, "#") {
			continue
		}

		if strings.HasPrefix(t, "include") {
			libs, err := filepath.Glob(t[8:])
			if err != nil {
				continue
			}
			for _, l := range libs {
				readLdSoConf(l, libpath)
			}
		} else {
			libpath = append(libpath, t)
		}
	}
}

func init() {
	deps = make(map[string]DepsInfo)
	deflib = []string{"/lib/", "/usr/lib/"}
	envlib = os.Getenv("LD_LIBRARY_PATH")
	readLdSoConf("/etc/ld.so.conf", conflib)

	flag.BoolVar(&verbose, "v", false, "Show binary info")
	flag.BoolVar(&showPath, "p", false, "Show library path")
	flag.BoolVar(&showTui, "tui", false, "Show it with TUI")
}

// search shared libraries as described in `man ld.so(8)`
func findLib(name string, parent *DepsNode) string {
	if strings.Contains(name, "/") {
		return name
	}

	// check DT_RPATH attribute
	if parent != nil {
		info := deps[parent.name]
		for _, dyn := range info.dyns {
			if dyn.tag != elf.DT_RPATH {
				continue
			}

			fullpath := path.Join(dyn.val.(string), name)
			if _, err := os.Stat(fullpath); err == nil {
				return fullpath
			}
		}
	}

	// check LD_LIBRARY_PATH environ
	for _, libpath := range strings.Split(envlib, ":") {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}

	// check DT_RUNPATH attribute
	if parent != nil {
		info := deps[parent.name]
		for _, dyn := range info.dyns {
			if dyn.tag != elf.DT_RUNPATH {
				continue
			}

			fullpath := path.Join(dyn.val.(string), name)
			if _, err := os.Stat(fullpath); err == nil {
				return fullpath
			}
		}
	}

	// check libraries in /etc/ld.so.conf
	for _, libpath := range conflib {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}

	// check default library directories
	for _, libpath := range deflib {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}
	return ""
}

func realPath(pathname string) string {
	if pathname == "" {
		return ""
	}

	relpath, _ := filepath.EvalSymlinks(pathname)
	abspath, _ := filepath.Abs(relpath)

	return abspath
}

func readElfString(strtab []byte, i uint64) string {
	var len uint64

	for len = 0; strtab[i+len] != '\x00'; len++ {
		continue
	}

	return string(strtab[i : i+len])
}

func readDynamic(f *elf.File, info *DepsInfo) {
	var i, count uint

	dyn := f.Section(".dynamic")
	data, err := dyn.Data()
	if err != nil {
		return
	}
	str := f.Section(".dynstr")
	stab, err := str.Data()
	if err != nil {
		return
	}

	count = uint(dyn.Size / dyn.Entsize)
	for i = 0; i < count; i++ {
		var tag, val uint64

		if f.Class == elf.ELFCLASS64 {
			tag = f.ByteOrder.Uint64(data[(i*2+0)*8 : (i*2+1)*8])
			val = f.ByteOrder.Uint64(data[(i*2+1)*8 : (i*2+2)*8])
		} else {
			tag = uint64(f.ByteOrder.Uint32(data[(i*2+0)*4 : (i*2+1)*4]))
			val = uint64(f.ByteOrder.Uint32(data[(i*2+1)*4 : (i*2+2)*4]))
		}

		dtag := elf.DynTag(tag)
		switch dtag {
		case elf.DT_NEEDED:
			fallthrough
		case elf.DT_RPATH:
			fallthrough
		case elf.DT_RUNPATH:
			fallthrough
		case elf.DT_SONAME:
			sval := readElfString(stab, val)
			info.dyns = append(info.dyns, DynInfo{dtag, sval})
			break
		default:
			info.dyns = append(info.dyns, DynInfo{dtag, val})
			break
		}
	}
}

func processDep(dep *DepsNode) {
	// skip duplicate libraries
	if _, ok := deps[dep.name]; ok {
		return
	}

	info := DepsInfo{path: realPath(findLib(dep.name, dep.parent))}

	if dep.parent == nil {
		info.path = realPath(flag.Args()[0])
	}

	f, err := elf.Open(info.path)
	if err != nil {
		fmt.Printf("%v: %s (%s)\n", err, info.path, dep.name)
		os.Exit(1)
	}
	defer f.Close()

	info.mach = f.Machine
	info.bits = f.Class
	info.kind = f.Type
	info.abi = f.OSABI
	info.ver = f.ABIVersion
	info.endian = f.ByteOrder

	info.prog = f.Progs

	readDynamic(f, &info)

	libs, err := f.ImportedLibraries()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	dsym, err := f.ImportedSymbols()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	syms, err := f.Symbols()
	if err == nil {
		info.syms = syms
	}

	info.libs = libs
	info.dsym = dsym

	var L []*DepsNode
	for _, soname := range libs {
		N := new(DepsNode)
		N.name = soname
		N.parent = dep
		N.depth = dep.depth + 1

		L = append(L, N)
		dep.child = append(dep.child, N)
	}

	deps_list = append(L, deps_list...)
	deps[dep.name] = info
}

func printDepTree(n *DepsNode, f *elf.File) {
	for i := 0; i < n.depth; i++ {
		fmt.Printf("   ")
	}

	if showPath {
		fmt.Printf("%s  => %s\n", n.name, deps[n.name].path)
	} else {
		fmt.Println(n.name)
	}

	for _, v := range n.child {
		printDepTree(v, f)
	}

	if verbose && n.parent == nil {
		showDetails(f, deps[n.name].path)
	}
}

func showDetails(f *elf.File, pathname string) {
	s := f.Section(".interp")
	if s == nil {
		fmt.Printf("static linked executable: %s\n", pathname)
		os.Exit(1)
	}
	interp, err := s.Data()
	if err != nil {
		fmt.Printf("%v: %s\n", err, pathname)
		os.Exit(1)
	}

	di_deps, err := f.ImportedLibraries()
	if err != nil {
		fmt.Printf("imported libraries: %v\n", err)
		os.Exit(1)
	}

	fmt.Println()
	fmt.Printf("%s: %s\n", path.Base(pathname), realPath(pathname))
	fmt.Printf("  type:                     %s  (%s / %s / %s)\n",
		f.Type, f.Machine, f.Class, f.ByteOrder)
	fmt.Printf("  interpreter:              %s\n", string(interp))
	fmt.Printf("  total dependency:         %d\n", len(deps)-1) // exclude itself
	fmt.Printf("  direct dependency:        %d\n", len(di_deps))
}

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) < 1 {
		fmt.Println("Usage: elftree [<options>] <executable>")
		os.Exit(1)
	}

	pathname := args[0]
	f, err := elf.Open(pathname)
	if err != nil {
		fmt.Printf("%v: %s\n", err, pathname)
		os.Exit(1)
	}
	defer f.Close()

	deps_root = new(DepsNode)
	deps_root.name = path.Base(pathname)

	deps_list = append(deps_list, deps_root)
	for len(deps_list) > 0 {
		// pop first element
		dep := deps_list[0]
		deps_list = deps_list[1:]

		processDep(dep)
	}

	if showTui {
		ShowWithTUI(deps_root)
	} else {
		printDepTree(deps_root, f)
	}
}
<commit_msg>Add {/usr,}/lib64 as default directories<commit_after>/* ELF tree - Tree viewer for ELF library dependency */
package main

import (
	"bufio"
	"debug/elf"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type DepsNode struct {
	name   string
	parent *DepsNode
	child  []*DepsNode
	depth  int
}

type DynInfo struct {
	tag elf.DynTag
	val interface{}
}

type DepsInfo struct {
	path   string
	mach   elf.Machine
	bits   elf.Class
	endian binary.ByteOrder
	kind   elf.Type
	abi    elf.OSABI
	ver    uint8

	libs []string
	dsym []elf.ImportedSymbol
	syms []elf.Symbol
	prog []*elf.Prog
	dyns []DynInfo
}

var (
	deps      map[string]DepsInfo
	deps_list []*DepsNode
	deps_root *DepsNode
	deflib    []string
	envlib    string
	conflib   []string
)

// command-line options
var (
	verbose  bool
	showPath bool
	showTui  bool
)

func readLdSoConf(name string, libpath []string) {
	f, err := os.Open(name)
	if err != nil {
		return
	}
	defer f.Close()

	s := bufio.NewScanner(f)
	for s.Scan() {
		t := s.Text()

		if len(strings.TrimSpace(t)) == 0 {
			continue
		}
		if strings.HasPrefix(t, "#") {
			continue
		}

		if strings.HasPrefix(t, "include") {
			libs, err := filepath.Glob(t[8:])
			if err != nil {
				continue
			}
			for _, l := range libs {
				readLdSoConf(l, libpath)
			}
		} else {
			libpath = append(libpath, t)
		}
	}
}

func init() {
	deps = make(map[string]DepsInfo)
	deflib = []string{"/lib/", "/usr/lib/", "/lib64", "/usr/lib64"}
	envlib = os.Getenv("LD_LIBRARY_PATH")
	readLdSoConf("/etc/ld.so.conf", conflib)

	flag.BoolVar(&verbose, "v", false, "Show binary info")
	flag.BoolVar(&showPath, "p", false, "Show library path")
	flag.BoolVar(&showTui, "tui", false, "Show it with TUI")
}

// search shared libraries as described in `man ld.so(8)`
func findLib(name string, parent *DepsNode) string {
	if strings.Contains(name, "/") {
		return name
	}

	// check DT_RPATH attribute
	if parent != nil {
		info := deps[parent.name]
		for _, dyn := range info.dyns {
			if dyn.tag != elf.DT_RPATH {
				continue
			}

			fullpath := path.Join(dyn.val.(string), name)
			if _, err := os.Stat(fullpath); err == nil {
				return fullpath
			}
		}
	}

	// check LD_LIBRARY_PATH environ
	for _, libpath := range strings.Split(envlib, ":") {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}

	// check DT_RUNPATH attribute
	if parent != nil {
		info := deps[parent.name]
		for _, dyn := range info.dyns {
			if dyn.tag != elf.DT_RUNPATH {
				continue
			}

			fullpath := path.Join(dyn.val.(string), name)
			if _, err := os.Stat(fullpath); err == nil {
				return fullpath
			}
		}
	}

	// check libraries in /etc/ld.so.conf
	for _, libpath := range conflib {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}

	// check default library directories
	for _, libpath := range deflib {
		fullpath := path.Join(libpath, name)
		if _, err := os.Stat(fullpath); err == nil {
			return fullpath
		}
	}
	return ""
}

func realPath(pathname string) string {
	if pathname == "" {
		return ""
	}

	relpath, _ := filepath.EvalSymlinks(pathname)
	abspath, _ := filepath.Abs(relpath)

	return abspath
}

func readElfString(strtab []byte, i uint64) string {
	var len uint64

	for len = 0; strtab[i+len] != '\x00'; len++ {
		continue
	}

	return string(strtab[i : i+len])
}

func readDynamic(f *elf.File, info *DepsInfo) {
	var i, count uint

	dyn := f.Section(".dynamic")
	data, err := dyn.Data()
	if err != nil {
		return
	}
	str := f.Section(".dynstr")
	stab, err := str.Data()
	if err != nil {
		return
	}

	count = uint(dyn.Size / dyn.Entsize)
	for i = 0; i < count; i++ {
		var tag, val uint64

		if f.Class == elf.ELFCLASS64 {
			tag = f.ByteOrder.Uint64(data[(i*2+0)*8 : (i*2+1)*8])
			val = f.ByteOrder.Uint64(data[(i*2+1)*8 : (i*2+2)*8])
		} else {
			tag = uint64(f.ByteOrder.Uint32(data[(i*2+0)*4 : (i*2+1)*4]))
			val = uint64(f.ByteOrder.Uint32(data[(i*2+1)*4 : (i*2+2)*4]))
		}

		dtag := elf.DynTag(tag)
		switch dtag {
		case elf.DT_NEEDED:
			fallthrough
		case elf.DT_RPATH:
			fallthrough
		case elf.DT_RUNPATH:
			fallthrough
		case elf.DT_SONAME:
			sval := readElfString(stab, val)
			info.dyns = append(info.dyns, DynInfo{dtag, sval})
			break
		default:
			info.dyns = append(info.dyns, DynInfo{dtag, val})
			break
		}
	}
}

func processDep(dep *DepsNode) {
	// skip duplicate libraries
	if _, ok := deps[dep.name]; ok {
		return
	}

	info := DepsInfo{path: realPath(findLib(dep.name, dep.parent))}

	if dep.parent == nil {
		info.path = realPath(flag.Args()[0])
	}

	f, err := elf.Open(info.path)
	if err != nil {
		fmt.Printf("%v: %s (%s)\n", err, info.path, dep.name)
		os.Exit(1)
	}
	defer f.Close()

	info.mach = f.Machine
	info.bits = f.Class
	info.kind = f.Type
	info.abi = f.OSABI
	info.ver = f.ABIVersion
	info.endian = f.ByteOrder

	info.prog = f.Progs

	readDynamic(f, &info)

	libs, err := f.ImportedLibraries()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	dsym, err := f.ImportedSymbols()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	syms, err := f.Symbols()
	if err == nil {
		info.syms = syms
	}

	info.libs = libs
	info.dsym = dsym

	var L []*DepsNode
	for _, soname := range libs {
		N := new(DepsNode)
		N.name = soname
		N.parent = dep
		N.depth = dep.depth + 1

		L = append(L, N)
		dep.child = append(dep.child, N)
	}

	deps_list = append(L, deps_list...)
	deps[dep.name] = info
}

func printDepTree(n *DepsNode, f *elf.File) {
	for i := 0; i < n.depth; i++ {
		fmt.Printf("   ")
	}

	if showPath {
		fmt.Printf("%s  => %s\n", n.name, deps[n.name].path)
	} else {
		fmt.Println(n.name)
	}

	for _, v := range n.child {
		printDepTree(v, f)
	}

	if verbose && n.parent == nil {
		showDetails(f, deps[n.name].path)
	}
}

func showDetails(f *elf.File, pathname string) {
	s := f.Section(".interp")
	if s == nil {
		fmt.Printf("static linked executable: %s\n", pathname)
		os.Exit(1)
	}
	interp, err := s.Data()
	if err != nil {
		fmt.Printf("%v: %s\n", err, pathname)
		os.Exit(1)
	}

	di_deps, err := f.ImportedLibraries()
	if err != nil {
		fmt.Printf("imported libraries: %v\n", err)
		os.Exit(1)
	}

	fmt.Println()
	fmt.Printf("%s: %s\n", path.Base(pathname), realPath(pathname))
	fmt.Printf("  type:                     %s  (%s / %s / %s)\n",
		f.Type, f.Machine, f.Class, f.ByteOrder)
	fmt.Printf("  interpreter:              %s\n", string(interp))
	fmt.Printf("  total dependency:         %d\n", len(deps)-1) // exclude itself
	fmt.Printf("  direct dependency:        %d\n", len(di_deps))
}

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) < 1 {
		fmt.Println("Usage: elftree [<options>] <executable>")
		os.Exit(1)
	}

	pathname := args[0]
	f, err := elf.Open(pathname)
	if err != nil {
		fmt.Printf("%v: %s\n", err, pathname)
		os.Exit(1)
	}
	defer f.Close()

	deps_root = new(DepsNode)
	deps_root.name = path.Base(pathname)

	deps_list = append(deps_list, deps_root)
	for len(deps_list) > 0 {
		// pop first element
		dep := deps_list[0]
		deps_list = deps_list[1:]

		processDep(dep)
	}

	if showTui {
		ShowWithTUI(deps_root)
	} else {
		printDepTree(deps_root, f)
	}
}
<|endoftext|>


<commit_before>package main

import (
	// Standard library
	"fmt"
	"os"
	"os/signal"
	"strings"

	// Internal packages
	"github.com/deuill/sigil/srv"

	// External packages
	"github.com/rakyll/globalconf"
)

func setup(name string) error {
	var (
		conf *globalconf.GlobalConf
		file string
		err  error
	)

	// Check for environment variable override, and return error if override is set but file is
	// inaccessible. Otherwise, try for the default global location (in the "/etc" directory).
	if file := os.Getenv(strings.ToUpper(name) + "_CONFIG"); file != "" {
		if _, err = os.Stat(file); err != nil {
			return err
		}
	} else {
		file = "/etc/" + name + "/" + name + ".conf"
		if _, err = os.Stat(file); err != nil {
			file = ""
		}
	}

	// Load from specific configuration file if set, or use local configuration file as a fallback.
	if file != "" {
		if conf, err = globalconf.NewWithOptions(&globalconf.Options{file, ""}); err != nil {
			return err
		}
	} else if conf, err = globalconf.New(name); err != nil {
		return err
	}

	conf.EnvPrefix = strings.ToUpper(name) + "_"
	conf.ParseAll()

	return nil
}

func main() {
	var err error

	if err = setup("sigil"); err != nil {
		fmt.Println("Error loading configuration:", err)
		os.Exit(1)
	}

	fmt.Print("Starting server... ")

	// Initialize HTTP and attached services.
	err = srv.Init()
	if err != nil {
		fmt.Printf("error initializing services:\n%s\n", err)
		os.Exit(1)
	}

	fmt.Println("done.")

	// Listen for and terminate Sigil on SIGKILL or SIGINT signals.
	sigStop := make(chan os.Signal)
	signal.Notify(sigStop, os.Interrupt, os.Kill)

	select {
	case <-sigStop:
		fmt.Println("\rShutting down server...")
	}

	os.Exit(0)
}
<commit_msg>Fix issue reported by `go vet`.<commit_after>package main

import (
	// Standard library
	"fmt"
	"os"
	"os/signal"
	"strings"

	// Internal packages
	"github.com/deuill/sigil/srv"

	// External packages
	"github.com/rakyll/globalconf"
)

func setup(name string) error {
	var (
		conf *globalconf.GlobalConf
		file string
		err  error
	)

	// Check for environment variable override, and return error if override is set but file is
	// inaccessible. Otherwise, try for the default global location (in the "/etc" directory).
	if file := os.Getenv(strings.ToUpper(name) + "_CONFIG"); file != "" {
		if _, err = os.Stat(file); err != nil {
			return err
		}
	} else {
		file = "/etc/" + name + "/" + name + ".conf"
		if _, err = os.Stat(file); err != nil {
			file = ""
		}
	}

	// Load from specific configuration file if set, or use local configuration file as a fallback.
	if file != "" {
		options := &globalconf.Options{Filename: file, EnvPrefix: ""}
		if conf, err = globalconf.NewWithOptions(options); err != nil {
			return err
		}
	} else if conf, err = globalconf.New(name); err != nil {
		return err
	}

	conf.EnvPrefix = strings.ToUpper(name) + "_"
	conf.ParseAll()

	return nil
}

func main() {
	var err error

	if err = setup("sigil"); err != nil {
		fmt.Println("Error loading configuration:", err)
		os.Exit(1)
	}

	fmt.Print("Starting server... ")

	// Initialize HTTP and attached services.
	err = srv.Init()
	if err != nil {
		fmt.Printf("error initializing services:\n%s\n", err)
		os.Exit(1)
	}

	fmt.Println("done.")

	// Listen for and terminate Sigil on SIGKILL or SIGINT signals.
	sigStop := make(chan os.Signal)
	signal.Notify(sigStop, os.Interrupt, os.Kill)

	select {
	case <-sigStop:
		fmt.Println("\rShutting down server...")
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/bradfitz/http2"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/pat"
	"github.com/gorilla/securecookie"
	"github.com/unrolled/secure"
	"gitlab.com/mattbostock/timeoff/handler"
	"gitlab.com/mattbostock/timeoff/middleware/negroni_logrus"
	"gitlab.com/mattbostock/timeoff/middleware/sessions"
	"gitlab.com/mattbostock/timeoff/model"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/github"
)

const (
	assetsPath  = "assets"
	defaultAddr = "localhost:3000"
	sessionName = "timeoff_session"
)

var (
	config = &struct {
		addr               string
		allowedHosts       []string
		cookieHashKey      []byte
		debug              bool
		gitHubClientID     string
		gitHubClientSecret string
		tlsCert            string
		tlsKey             string
	}{
		addr:               os.Getenv("ADDR"),
		cookieHashKey:      []byte(os.Getenv("COOKIE_KEY")),
		debug:              os.Getenv("DEBUG") != "",
		gitHubClientID:     os.Getenv("GITHUB_CLIENT_ID"),
		gitHubClientSecret: os.Getenv("GITHUB_CLIENT_SECRET"),
		tlsCert:            os.Getenv("TLS_CERT"),
		tlsKey:             os.Getenv("TLS_KEY"),
	}

	mux     = pat.New()
	log     = logrus.New()
	version = ""
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	if config.debug {
		log.Level = logrus.DebugLevel
	}

	if config.addr == "" {
		config.addr = defaultAddr
	}
}

func main() {
	if version == "" {
		log.Fatalln(errMakeFileNotUsed)
	}

	if os.Getenv("ALLOWED_HOSTS") == "" {
		config.allowedHosts = nil
	} else {
		config.allowedHosts = strings.Split(os.Getenv("ALLOWED_HOSTS"), ",")
	}

	if len(config.cookieHashKey) == 0 {
		log.Warningln(errNoCookieHashKey)
		config.cookieHashKey = securecookie.GenerateRandomKey(32)
	}
	if len(config.cookieHashKey) != 32 {
		// additonal check as securecookie.GenerateRandomKey() does not return errors
		log.Fatalf(errCookieHashKeyWrongLength, len(config.cookieHashKey))
	}

	if config.gitHubClientID == "" || config.gitHubClientSecret == "" {
		log.Fatalf(errNoGitHubCredentials)
	}

	handler.SetVersion(version)

	handler.SetOauthConfig(&oauth2.Config{
		ClientID:     config.gitHubClientID,
		ClientSecret: config.gitHubClientSecret,
		Endpoint:     github.Endpoint,
		Scopes:       []string{"user:email"},
	})

	model.SetLogger(log)
	model.InitDB("sqlite3", "sqlite.db")

	sessions.SetLogger(log)
	handler.SetLogger(log)

	sessionManager := sessions.New(sessionName, config.cookieHashKey)
	handler.SetSessionManager(sessionManager)

	secureOpts := secure.Options{
		AllowedHosts:          config.allowedHosts,
		BrowserXssFilter:      true,
		ContentSecurityPolicy: "default-src 'self'; img-src 'self' data:",
		FrameDeny:             true,
		STSIncludeSubdomains:  true,
		STSSeconds:            365 * 24 * 60 * 60,
	}
	secureMiddleware := secure.New(secureOpts)

	n := negroni.New()
	n.Use(negroniLogrus.New(log)) // logger must be first middleware
	n.Use(negroni.NewRecovery())
	n.Use(negroni.HandlerFunc(secureMiddleware.HandlerFuncWithNext))
	n.Use(negroni.NewStatic(http.Dir(assetsPath)))
	n.UseHandler(sessionManager)
	n.UseHandler(mux)
	registerRoutes()

	log.Infof("Listening on %s", config.addr)

	if config.allowedHosts != nil {
		log.Infof("Allowed HTTP hosts: %q", config.allowedHosts)
	} else {
		log.Warningln(errAllHostsEnabled)
	}

	c := &tls.Config{MinVersion: tls.VersionTLS10} // disable SSLv3, prevent POODLE attack
	s := &http.Server{Addr: config.addr, Handler: n, TLSConfig: c}

	if config.tlsCert == "" && config.tlsKey == "" {
		log.Warningln(errNoTLSCertificate)
		log.Fatal(s.ListenAndServe())
	} else {
		http2.ConfigureServer(s, nil)
		log.Infoln("TLS-only; HTTP/2 enabled")
		log.Fatal(s.ListenAndServeTLS(config.tlsCert, config.tlsKey))
	}
}

const (
	errAllHostsEnabled = "Accepting connections for all HTTP hosts. Consider setting the ALLOWED_HOSTS environment variable."

	errCookieHashKeyWrongLength = "COOKIE_KEY environment variable must be 32 characters long. Length provided: %d"

	errMakeFileNotUsed = "Makefile was not used when compiling binary, run 'make' to re-compile"

	errNoCookieHashKey = "No cookie hash key supplied. You should set the COOKIE_KEY " +
		"environment variable in a production environment. Falling back to use a temporary key " +
		"which will persist only for the current running process."

	errNoGitHubCredentials = "No GitHub Oauth credentials supplied. Set both GITHUB_CLIENT_ID and " +
		"GITHUB_CLIENT_SECRET environment variables."

	errNoTLSCertificate = "No TLS certficiate supplied. Consider setting TLS_CERT " +
		"and TLS_KEY environment variables to enable TLS. TimeOff will not work unless you " +
		"are using TLS upstream."
)
<commit_msg>Add -version command-line flag<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/bradfitz/http2"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/pat"
	"github.com/gorilla/securecookie"
	"github.com/unrolled/secure"
	"gitlab.com/mattbostock/timeoff/handler"
	"gitlab.com/mattbostock/timeoff/middleware/negroni_logrus"
	"gitlab.com/mattbostock/timeoff/middleware/sessions"
	"gitlab.com/mattbostock/timeoff/model"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/github"
)

const (
	assetsPath  = "assets"
	defaultAddr = "localhost:3000"
	sessionName = "timeoff_session"
)

var (
	config = &struct {
		addr               string
		allowedHosts       []string
		cookieHashKey      []byte
		debug              bool
		gitHubClientID     string
		gitHubClientSecret string
		tlsCert            string
		tlsKey             string
	}{
		addr:               os.Getenv("ADDR"),
		cookieHashKey:      []byte(os.Getenv("COOKIE_KEY")),
		debug:              os.Getenv("DEBUG") != "",
		gitHubClientID:     os.Getenv("GITHUB_CLIENT_ID"),
		gitHubClientSecret: os.Getenv("GITHUB_CLIENT_SECRET"),
		tlsCert:            os.Getenv("TLS_CERT"),
		tlsKey:             os.Getenv("TLS_KEY"),
	}

	mux     = pat.New()
	log     = logrus.New()
	version = ""
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	if config.debug {
		log.Level = logrus.DebugLevel
	}

	if config.addr == "" {
		config.addr = defaultAddr
	}
}

func main() {
	if version == "" {
		log.Fatalln(errMakeFileNotUsed)
	}

	v := flag.Bool("version", false, "prints current version")
	flag.Parse()

	if *v {
		fmt.Println(version)
		os.Exit(0)
	}

	if os.Getenv("ALLOWED_HOSTS") == "" {
		config.allowedHosts = nil
	} else {
		config.allowedHosts = strings.Split(os.Getenv("ALLOWED_HOSTS"), ",")
	}

	if len(config.cookieHashKey) == 0 {
		log.Warningln(errNoCookieHashKey)
		config.cookieHashKey = securecookie.GenerateRandomKey(32)
	}
	if len(config.cookieHashKey) != 32 {
		// additonal check as securecookie.GenerateRandomKey() does not return errors
		log.Fatalf(errCookieHashKeyWrongLength, len(config.cookieHashKey))
	}

	if config.gitHubClientID == "" || config.gitHubClientSecret == "" {
		log.Fatalf(errNoGitHubCredentials)
	}

	handler.SetVersion(version)

	handler.SetOauthConfig(&oauth2.Config{
		ClientID:     config.gitHubClientID,
		ClientSecret: config.gitHubClientSecret,
		Endpoint:     github.Endpoint,
		Scopes:       []string{"user:email"},
	})

	model.SetLogger(log)
	model.InitDB("sqlite3", "sqlite.db")

	sessions.SetLogger(log)
	handler.SetLogger(log)

	sessionManager := sessions.New(sessionName, config.cookieHashKey)
	handler.SetSessionManager(sessionManager)

	secureOpts := secure.Options{
		AllowedHosts:          config.allowedHosts,
		BrowserXssFilter:      true,
		ContentSecurityPolicy: "default-src 'self'; img-src 'self' data:",
		FrameDeny:             true,
		STSIncludeSubdomains:  true,
		STSSeconds:            365 * 24 * 60 * 60,
	}
	secureMiddleware := secure.New(secureOpts)

	n := negroni.New()
	n.Use(negroniLogrus.New(log)) // logger must be first middleware
	n.Use(negroni.NewRecovery())
	n.Use(negroni.HandlerFunc(secureMiddleware.HandlerFuncWithNext))
	n.Use(negroni.NewStatic(http.Dir(assetsPath)))
	n.UseHandler(sessionManager)
	n.UseHandler(mux)
	registerRoutes()

	log.Infof("Listening on %s", config.addr)

	if config.allowedHosts != nil {
		log.Infof("Allowed HTTP hosts: %q", config.allowedHosts)
	} else {
		log.Warningln(errAllHostsEnabled)
	}

	c := &tls.Config{MinVersion: tls.VersionTLS10} // disable SSLv3, prevent POODLE attack
	s := &http.Server{Addr: config.addr, Handler: n, TLSConfig: c}

	if config.tlsCert == "" && config.tlsKey == "" {
		log.Warningln(errNoTLSCertificate)
		log.Fatal(s.ListenAndServe())
	} else {
		http2.ConfigureServer(s, nil)
		log.Infoln("TLS-only; HTTP/2 enabled")
		log.Fatal(s.ListenAndServeTLS(config.tlsCert, config.tlsKey))
	}
}

const (
	errAllHostsEnabled = "Accepting connections for all HTTP hosts. Consider setting the ALLOWED_HOSTS environment variable."

	errCookieHashKeyWrongLength = "COOKIE_KEY environment variable must be 32 characters long. Length provided: %d"

	errMakeFileNotUsed = "Makefile was not used when compiling binary, run 'make' to re-compile"

	errNoCookieHashKey = "No cookie hash key supplied. You should set the COOKIE_KEY " +
		"environment variable in a production environment. Falling back to use a temporary key " +
		"which will persist only for the current running process."

	errNoGitHubCredentials = "No GitHub Oauth credentials supplied. Set both GITHUB_CLIENT_ID and " +
		"GITHUB_CLIENT_SECRET environment variables."

	errNoTLSCertificate = "No TLS certficiate supplied. Consider setting TLS_CERT " +
		"and TLS_KEY environment variables to enable TLS. TimeOff will not work unless you " +
		"are using TLS upstream."
)
<|endoftext|>


<commit_before>// Package hdf5 provides a reader and writer of HDF5 files.
//
// https://en.wikipedia.org/wiki/Hierarchical_Data_Format
package hdf5

// #cgo CFLAGS: -Ihdf5/install/include
// #cgo LDFLAGS: -lm -lz
//
// #include <stdlib.h>
// #include <string.h>
//
// #include <hdf5.h>
// #include <hdf5_hl.h>
import "C"

import (
	"errors"
	"unsafe"
)

// File represents a file.
type File struct {
	id C.hid_t
}

// Open opens a file for reading and writing.
func Open(path string) (*File, error) {
	const (
		// https://github.com/copies/hdf5/blob/master/src/H5Fpublic.h#L46
		F_ACC_RDWR = 0x0001
	)

	cpath := C.CString(path)
	defer C.free(unsafe.Pointer(cpath))

	id := C.H5Fopen(cpath, F_ACC_RDWR, C.H5P_DEFAULT)
	if id < 0 {
		return nil, errors.New("failed to open the file")
	}

	file := &File{
		id: id,
	}

	return file, nil
}

// Close closes the file.
func (f *File) Close() error {
	if err := C.H5Fclose(f.id); err != 0 {
		return errors.New("failed to close the file")
	}
	return nil
}
<commit_msg>Fix LDFLAGS<commit_after>// Package hdf5 provides a reader and writer of HDF5 files.
//
// https://en.wikipedia.org/wiki/Hierarchical_Data_Format
package hdf5

// #cgo CFLAGS: -Ihdf5/install/include
// #cgo LDFLAGS: -ldl -lm -lz
//
// #include <stdlib.h>
// #include <string.h>
//
// #include <hdf5.h>
// #include <hdf5_hl.h>
import "C"

import (
	"errors"
	"unsafe"
)

// File represents a file.
type File struct {
	id C.hid_t
}

// Open opens a file for reading and writing.
func Open(path string) (*File, error) {
	const (
		// https://github.com/copies/hdf5/blob/master/src/H5Fpublic.h#L46
		F_ACC_RDWR = 0x0001
	)

	cpath := C.CString(path)
	defer C.free(unsafe.Pointer(cpath))

	id := C.H5Fopen(cpath, F_ACC_RDWR, C.H5P_DEFAULT)
	if id < 0 {
		return nil, errors.New("failed to open the file")
	}

	file := &File{
		id: id,
	}

	return file, nil
}

// Close closes the file.
func (f *File) Close() error {
	if err := C.H5Fclose(f.id); err != 0 {
		return errors.New("failed to close the file")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"strings"

	"./schemas"
	"github.com/Shopify/sarama"
	"github.com/companieshouse/chs.go/avro"
	"github.com/companieshouse/chs.go/avro/schema"
	"github.com/companieshouse/chs.go/kafka/producer"
)

// Assigns all flags to variables
var (
	brokerPtr         = flag.String("broker", "", "Broker address")
	topicPtr          = flag.String("topic", "", "Topic name")
	schemaPtr         = flag.String("schema", "", "Schema name")
	schemaRegistryPtr = flag.String("schema-registry", "", "Schema Registry")
	partitionPtr      = flag.Int64("partition", 0, "Partition (default: 0)")
	offsetPtr         = flag.String("offset", "", "Offset number")
	jsonOutPtr        = flag.Int64("json-out", 0, "Print deserialized JSON message (default: 0)")
)

func main() {
	flag.Parse()

	// check if all mandatory flags have been provided
	validateFlags()

	// Print flags and parameters of the tool
	printFlags()

	// create offset array
	offsetArray := createOffsetArray(*offsetPtr)

	// create default config for sarama
	config := sarama.NewConfig()

	//create consumer
	consumer, err := sarama.NewConsumer([]string{*brokerPtr}, config)
	if err != nil {
		panic(err)
	}

	// create messages chan
	messages := make(chan *sarama.ConsumerMessage)
	go consumePartition(consumer, *topicPtr, int32(*partitionPtr), messages, offsetArray)

	// create signals chan
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	consumed := 0
ConsumerLoop:
	for {
		select {
		case msg := <-messages:

			schemaStruct := schemas.IdentifySchema(*topicPtr)
			var err error

			// Get schema from schema registry and use it to create avro consumer
			schema, err := schema.Get(*schemaRegistryPtr, *schemaPtr)
			consumerAvro := &avro.Schema{
				Definition: schema,
			}

			// Unmarshal message value and assign it to schemaStruct
			if err = consumerAvro.Unmarshal(msg.Value, schemaStruct); err != nil {
				fmt.Errorf("error unmarshalling avro: %v", err)
				os.Exit(1)
			}

			// create avro producer
			producerAvro := &avro.Schema{
				Definition: schema,
			}

			// Marshall schemaStruct that contains message value ready for republishing
			messageBytes, err := producerAvro.Marshal(schemaStruct)
			if err != nil {
				fmt.Errorf("error marshalling avro: %v", err)
				os.Exit(1)
			}

			// create producer message
			producerMessage := &producer.Message{
				Value: messageBytes,
				Topic: *topicPtr,
			}

			//create new producer
			p, err := producer.New(&producer.Config{Acks: &producer.WaitForAll, BrokerAddrs: []string{*brokerPtr}})
			if err != nil {
				fmt.Errorf("error creating producer: %v", err)
				os.Exit(1)
			}

			// republish message
			partition, offset, err := p.Send(producerMessage)
			if err != nil {
				fmt.Errorf("error republishing message to topic: %v", err)
				os.Exit(1)
			}

			fmt.Printf("Message republished to topic: %v using partition: %v offset: %v \n", *topicPtr, partition, offset)

			consumed++
		case <-signals:
			fmt.Println("break")
			break ConsumerLoop
		}
	}
}

// Validates the flags to make sure all mandatory flags have been supplied
// throw an error if not the case
func validateFlags() {
	flag.VisitAll(func(f *flag.Flag) {
		if string(f.Value.String()) == "" {
			fmt.Printf("Value not supplied for: %v \n", f.Name)
			os.Exit(1)
		}
	})
}

// Prints the flags and parameters of the tool
func printFlags() {
	fmt.Println("---------------------")
	fmt.Println("Parameters:")
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Printf("%v: %v\n", f.Name, f.Value)
	})
	fmt.Println("---------------------")
}

// The offset slice is passed into method and is iterated over and each message
// related to the offset in the topic/partition is outputted into the 'out' chan
func consumePartition(consumer sarama.Consumer, topic string, partition int32, out chan *sarama.ConsumerMessage, offsetArray []int64) {
	for offset := range offsetArray {
		partitionConsumer, err := consumer.ConsumePartition(topic, partition, int64(offset))
		if err != nil {
			panic(err)
		}

		select {
		case msg := <-partitionConsumer.Messages():
			partitionConsumer.Close()
			out <- msg
		}
	}
}

// This creates the offset array depending on whether a range or a single value
// is entered as an argument to the tool
func createOffsetArray(offset string) []int64 {
	arraySize := make([]int64, 0)
	if strings.ContainsAny(offset, "-") {
		slice := strings.Split(offset, "-")
		minRange, err := strconv.ParseInt(slice[0], 10, 64)
		maxRange, err := strconv.ParseInt(slice[1], 10, 64)
		if err != nil {
			panic(err)
		}

		index := 0
		for value := minRange; value <= maxRange; value++ {
			arraySize = append(arraySize, value)
			index++
		}
	} else {
		value, err := strconv.ParseInt(offset, 10, 64)
		if err != nil {
			panic(err)
		}
		arraySize = append(arraySize, value)
	}
	return arraySize
}
<commit_msg>Improved logging to output progress<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"strings"

	"./schemas"
	"github.com/Shopify/sarama"
	"github.com/companieshouse/chs.go/avro"
	"github.com/companieshouse/chs.go/avro/schema"
	"github.com/companieshouse/chs.go/kafka/producer"
)

// Assigns all flags to variables
var (
	brokerPtr         = flag.String("broker", "", "Broker address")
	topicPtr          = flag.String("topic", "", "Topic name")
	schemaPtr         = flag.String("schema", "", "Schema name")
	schemaRegistryPtr = flag.String("schema-registry", "", "Schema Registry")
	partitionPtr      = flag.Int64("partition", 0, "Partition (default: 0)")
	offsetPtr         = flag.String("offset", "", "Offset number")
	jsonOutPtr        = flag.Int64("json-out", 0, "Print deserialized JSON message (default: 0)")
)

func main() {
	flag.Parse()

	// check if all mandatory flags have been provided
	validateFlags()

	// Print flags and parameters of the tool
	printFlags()

	// create offset array
	offsetArray := createOffsetArray(*offsetPtr)

	// create default config for sarama
	config := sarama.NewConfig()

	//create consumer
	consumer, err := sarama.NewConsumer([]string{*brokerPtr}, config)
	if err != nil {
		panic(err)
	}

	// create messages chan
	messages := make(chan *sarama.ConsumerMessage)
	go consumePartition(consumer, *topicPtr, int32(*partitionPtr), messages, offsetArray)

	// create signals chan
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	consumed := 0
ConsumerLoop:
	for {
		select {
		case msg := <-messages:

			schemaStruct := schemas.IdentifySchema(*topicPtr)
			var err error

			// Get schema from schema registry and use it to create avro consumer
			fmt.Printf("Getting schema for topic: %v \n", *topicPtr)
			schema, err := schema.Get(*schemaRegistryPtr, *schemaPtr)
			if err != nil {
				fmt.Errorf("rror getting schema: %v", err)
				os.Exit(1)
			}
			fmt.Printf("Retrieved schema for topic: %v \n", *topicPtr)

			consumerAvro := &avro.Schema{
				Definition: schema,
			}

			// Unmarshal message value and assign it to schemaStruct
			fmt.Printf("Unmarshalling message for offset: %v \n", msg.Offset)
			if err = consumerAvro.Unmarshal(msg.Value, schemaStruct); err != nil {
				fmt.Errorf("error unmarshalling avro message for offset: %v, %v \n", msg.Offset, err)
				os.Exit(1)
			}
			fmt.Printf("Message successfully unmarshalled for offset: %v \n", msg.Offset)

			// create avro producer
			producerAvro := &avro.Schema{
				Definition: schema,
			}

			// Marshall schemaStruct that contains message value ready for republishing
			fmt.Printf("Marshalling message for offset: %v \n", msg.Offset)
			messageBytes, err := producerAvro.Marshal(schemaStruct)
			if err != nil {
				fmt.Errorf("error marshalling avro message for offset: %v, %v \n", msg.Offset, err)
				os.Exit(1)
			}
			fmt.Printf("Message successfully marshalled for offset: %v \n", msg.Offset)

			// create producer message
			producerMessage := &producer.Message{
				Value: messageBytes,
				Topic: *topicPtr,
			}

			//create new producer
			p, err := producer.New(&producer.Config{Acks: &producer.WaitForAll, BrokerAddrs: []string{*brokerPtr}})
			if err != nil {
				fmt.Errorf("error creating producer: %v", err)
				os.Exit(1)
			}

			// republish message
			fmt.Printf("Republishing message for offset: %v \n", msg.Offset)
			partition, offset, err := p.Send(producerMessage)
			if err != nil {
				fmt.Errorf("error republishing message with offset: %v, %v \n", msg.Offset, err)
				os.Exit(1)
			}
			fmt.Printf("Successfully republished message with offset %v to topic: %v using partition: %v new offset: %v \n", msg.Offset, *topicPtr, partition, offset)
			fmt.Println("---------------------")

			consumed++
		case <-signals:
			fmt.Println("break")
			break ConsumerLoop
		}
	}
}

// Validates the flags to make sure all mandatory flags have been supplied
// throw an error if not the case
func validateFlags() {
	flag.VisitAll(func(f *flag.Flag) {
		if string(f.Value.String()) == "" {
			fmt.Printf("Value not supplied for: %v \n", f.Name)
			os.Exit(1)
		}
	})
}

// Prints the flags and parameters of the tool
func printFlags() {
	fmt.Println("---------------------")
	fmt.Println("Parameters:")
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Printf("%v: %v\n", f.Name, f.Value)
	})
	fmt.Println("---------------------")
}

// The offset slice is passed into method and is iterated over and each message
// related to the offset in the topic/partition is outputted into the 'out' chan
func consumePartition(consumer sarama.Consumer, topic string, partition int32, out chan *sarama.ConsumerMessage, offsetArray []int64) {
	for _, offset := range offsetArray {
		partitionConsumer, err := consumer.ConsumePartition(topic, partition, int64(offset))
		if err != nil {
			panic(err)
		}

		select {
		case msg := <-partitionConsumer.Messages():
			partitionConsumer.Close()
			out <- msg
		}
	}
}

// This creates the offset array depending on whether a range or a single value
// is entered as an argument to the tool
func createOffsetArray(offset string) []int64 {
	arraySize := make([]int64, 0)
	if strings.ContainsAny(offset, "-") {
		slice := strings.Split(offset, "-")
		minRange, err := strconv.ParseInt(slice[0], 10, 64)
		maxRange, err := strconv.ParseInt(slice[1], 10, 64)
		if err != nil {
			panic(err)
		}

		index := 0
		for value := minRange; value <= maxRange; value++ {
			arraySize = append(arraySize, value)
			index++
		}
	} else {
		value, err := strconv.ParseInt(offset, 10, 64)
		if err != nil {
			panic(err)
		}
		arraySize = append(arraySize, value)
	}
	return arraySize
}
<|endoftext|>


<commit_before>package main

import (
    "flag"
    "fmt"
    "runtime"
    "strconv"
    "strings"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

const Version = "0.2.0"

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
    PgsqlProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis", "pgsql"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {
    defer RECOVER("decodePktEth exception")

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet too short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet too short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short to contain TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    data_offset := (tcphdr[12] >> 4) * 4

    if l2hlen+iphl+int(data_offset) > l2hlen+ip_length {
        DEBUG("ip", "data_offset pointing outside of packet")
    }

    packet.payload = pkt.Data[l2hlen+iphl+int(data_offset) : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func main() {

    configfile := flag.String("c", "packetbeat.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 1, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    topSpeed := flag.Bool("t", false, "Read packets as fast as possible, without sleeping")
    publishDisabled := flag.Bool("N", false, "Disable actual publishing for testing")
    printVersion := flag.Bool("version", false, "Print version and exit")

    flag.Parse()

    if *printVersion {
        fmt.Printf("Packetbat version %s (%s)\n", Version, runtime.GOARCH)
    }

    debugSelectors := []string{}
    logLevel := LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. Exiting.\n", *configfile, err)
        return
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(*publishDisabled); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsupported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil && !*topSpeed {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went backwards: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<commit_msg>Disable standard logging by default.<commit_after>package main

import (
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "runtime"
    "strconv"
    "strings"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

const Version = "0.2.0"

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
    PgsqlProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis", "pgsql"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {
    defer RECOVER("decodePktEth exception")

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet too short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet too short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short to contain TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    data_offset := (tcphdr[12] >> 4) * 4

    if l2hlen+iphl+int(data_offset) > l2hlen+ip_length {
        DEBUG("ip", "data_offset pointing outside of packet")
    }

    packet.payload = pkt.Data[l2hlen+iphl+int(data_offset) : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func main() {

    configfile := flag.String("c", "packetbeat.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 1, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    topSpeed := flag.Bool("t", false, "Read packets as fast as possible, without sleeping")
    publishDisabled := flag.Bool("N", false, "Disable actual publishing for testing")
    printVersion := flag.Bool("version", false, "Print version and exit")

    flag.Parse()

    if *printVersion {
        fmt.Printf("Packetbat version %s (%s)\n", Version, runtime.GOARCH)
    }

    debugSelectors := []string{}
    logLevel := LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. Exiting.\n", *configfile, err)
        return
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if !IS_DEBUG("stdlog") {
        // disable standard logging by default
        log.SetOutput(ioutil.Discard)
    }

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(*publishDisabled); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsupported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil && !*topSpeed {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went backwards: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<|endoftext|>


<commit_before>/*
 * The MIT License (MIT)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package main

import (
	"github.com/gorilla/websocket"
	"html/template"
	"log"
	"net/http"
	"github.com/rvelhote/dnspropagation"
)

type WebsocketRequest struct {
    Domain     string `json:"domain"`
    RecordType string `json:"type"`
}

func index(w http.ResponseWriter, req *http.Request) {
	w.Header().Add("Content-Type", "text/html")
	t, _ := template.New("index.html").ParseFiles("templates/index.html")
	t.Execute(w, map[string]string{"Title": "Check DNS Propagation Worldwide"})
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:    4096,
	WriteBufferSize:   4096,
	EnableCompression: true,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func query(w http.ResponseWriter, req *http.Request, configuration []dnspropagation.Server) {
	conn, _ := upgrader.Upgrade(w, req, nil)

	websocketreq := WebsocketRequest{}
	conn.ReadJSON(&websocketreq)

	if len(websocketreq.Domain) == 0 {
		log.Print("Empty domain")
	}

	if len(websocketreq.RecordType) == 0 || dnspropagation.RecordTypes[websocketreq.RecordType] == 0 {
		log.Print("Invalid DNS record specified")
	}

	sem := make(chan dnspropagation.Response, len(configuration))

	for _, server := range configuration {
		go func(server dnspropagation.Server) {
			request := dnspropagation.DnsQuery{ Domain: websocketreq.Domain, Record: websocketreq.RecordType, Server: server }
			sem <- request.GetResponse()
		}(server)
	}

	for i := 0; i < len(configuration); i++ {
		conn.WriteJSON(<-sem)
	}

	conn.Close()
}

func main() {
	servers, _ := dnspropagation.LoadConfiguration("conf/servers.json")

	fs := http.FileServer(http.Dir("assets"))
	http.Handle("/assets/", http.StripPrefix("/assets/", fs))

	http.HandleFunc("/", index)

	http.HandleFunc("/api/v1/query", func(w http.ResponseWriter, req *http.Request) {
		query(w, req, servers)
	})

	http.ListenAndServe(":8080", nil)
}
<commit_msg>Use defer to close the WebSocket connection<commit_after>/*
 * The MIT License (MIT)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package main

import (
	"github.com/gorilla/websocket"
	"html/template"
	"log"
	"net/http"
	"github.com/rvelhote/dnspropagation"
)

type WebsocketRequest struct {
    Domain     string `json:"domain"`
    RecordType string `json:"type"`
}

func index(w http.ResponseWriter, req *http.Request) {
	w.Header().Add("Content-Type", "text/html")
	t, _ := template.New("index.html").ParseFiles("templates/index.html")
	t.Execute(w, map[string]string{"Title": "Check DNS Propagation Worldwide"})
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:    4096,
	WriteBufferSize:   4096,
	EnableCompression: true,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func query(w http.ResponseWriter, req *http.Request, configuration []dnspropagation.Server) {
	conn, _ := upgrader.Upgrade(w, req, nil)

	websocketreq := WebsocketRequest{}
	conn.ReadJSON(&websocketreq)

	defer conn.Close()

	if len(websocketreq.Domain) == 0 {
		log.Print("Empty domain")
	}

	if len(websocketreq.RecordType) == 0 || dnspropagation.RecordTypes[websocketreq.RecordType] == 0 {
		log.Print("Invalid DNS record specified")
	}

	sem := make(chan dnspropagation.Response, len(configuration))

	for _, server := range configuration {
		go func(server dnspropagation.Server) {
			request := dnspropagation.DnsQuery{ Domain: websocketreq.Domain, Record: websocketreq.RecordType, Server: server }
			sem <- request.GetResponse()
		}(server)
	}

	for i := 0; i < len(configuration); i++ {
		conn.WriteJSON(<-sem)
	}
}

func main() {
	servers, _ := dnspropagation.LoadConfiguration("conf/servers.json")

	fs := http.FileServer(http.Dir("assets"))
	http.Handle("/assets/", http.StripPrefix("/assets/", fs))

	http.HandleFunc("/", index)

	http.HandleFunc("/api/v1/query", func(w http.ResponseWriter, req *http.Request) {
		query(w, req, servers)
	})

	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>//go:generate esc -o ./app/filesystem.go -pkg app static
//go:generate go-bindata -o ./app/template-data.go -pkg app templates
package main

import (
	"log"
	"net/http"

	"github.com/kyleterry/sufr/app"
	"github.com/kyleterry/sufr/config"
)

func main() {
	config.New()
	log.Println("Welcome to SUFR")

	sufrApp := app.New()

	log.Printf("Listening on http://%s", config.ApplicationBind)
	http.ListenAndServe(config.ApplicationBind, sufrApp)
}
<commit_msg>handle errors for listen and serve<commit_after>//go:generate esc -o ./app/filesystem.go -pkg app static
//go:generate go-bindata -o ./app/template-data.go -pkg app templates
package main

import (
	"log"
	"net/http"

	"github.com/kyleterry/sufr/app"
	"github.com/kyleterry/sufr/config"
)

func main() {
	config.New()
	log.Println("Welcome to SUFR")

	sufrApp := app.New()

	log.Printf("Listening on http://%s", config.ApplicationBind)
	if err := http.ListenAndServe(config.ApplicationBind, sufrApp); err != nil {
		panic(err)
	}
}
<|endoftext|>