<commit_before>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	flagTimeout  = flag.Duration("t", 100*time.Millisecond, "upstream query timeout")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntTimedout  = expvar.NewInt("statsTimedout")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	cntRules     = expvar.NewInt("statsRules")
	answer       = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*net.UDPAddr
	blocked  map[string]bool
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - servers' bind address, e.g. 127.0.0.1\n"+
			"list.txt - text file with domains to block\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*net.UDPAddr, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerUpstreamDNS()
	go runServerLocalDNS()

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

forever:
	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGUSR1:
				log.Println("SIGUSR1 received, reloading rules")
			default:
				log.Println("Signal received, stopping")
				break forever
			}
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}
	log.Printf("DNS: Parsed %d entries from list\n", counter)
	cntRules.Set(int64(counter))

	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, nil)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, nil)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if to, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], to)
			if err != nil {
				log.Println("DNS ERROR:", err)
				cntErrors.Add(1)
				continue
			}
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
			cntRelayed.Add(1)
		}
	}
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var block bool

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	count := uint8(msg[5]) // question counter
	offset := 12           // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

	for {
		length := int8(msg[offset])
		if length == 0 {
			break
		}
		offset++
		domain.WriteString(string(msg[offset : offset+int(length)]))
		domain.WriteString(".")
		offset += int(length)
	}
	host := domain.String()
	testHost := host
	parts := strings.Split(testHost, ".")
	try := 1
	for {
		if _, ok := blocked[testHost]; ok {
			block = true
			break
		}
		parts = parts[1:]
		if len(parts) < 3 {
			break
		}
		testHost = strings.Join(parts, ".")
		try++
	}

	if block {
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}
		cntBlocked.Add(1)

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		msg = append(msg, msg[12:12+1+len(host)]...) // domain
		msg = append(msg, answer...)                 // payload
		_, err := proxy.WriteTo(msg, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = from
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			delete(queries, id)
			return
		}
		go func(queryID int) {
			time.Sleep(*flagTimeout)
			if queryFrom, ok := queries[queryID]; ok {
				fmt.Printf("DNS WARN: Query id %d from %s timed out\n", queryID, queryFrom)
				cntTimedout.Add(1)
				delete(queries, queryID)
			}
			return
		}(id)
	}

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<commit_msg>Verbose time out logging<commit_after>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

type query struct {
	Host string
	From *net.UDPAddr
}

func (q *query) String() string {
	return fmt.Sprintf("from %s about %s", q.From, q.Host)
}

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	flagTimeout  = flag.Duration("t", 100*time.Millisecond, "upstream query timeout")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntTimedout  = expvar.NewInt("statsTimedout")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	cntRules     = expvar.NewInt("statsRules")
	answer       = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*query
	blocked  map[string]bool
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - servers' bind address, e.g. 127.0.0.1\n"+
			"list.txt - text file with domains to block\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*query, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerUpstreamDNS()
	go runServerLocalDNS()

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

forever:
	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGUSR1:
				log.Println("SIGUSR1 received, reloading rules")
			default:
				log.Println("Signal received, stopping")
				break forever
			}
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}
	log.Printf("DNS: Parsed %d entries from list\n", counter)
	cntRules.Set(int64(counter))

	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, nil)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, nil)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if query, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], query.From)
			if err != nil {
				log.Printf("DNS ERROR: Query id %d %s %s", id, query, err)
				cntErrors.Add(1)
				continue
			}
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
			cntRelayed.Add(1)
		}
	}
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var block bool

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	count := uint8(msg[5]) // question counter
	offset := 12           // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

	for {
		length := int8(msg[offset])
		if length == 0 {
			break
		}
		offset++
		domain.WriteString(string(msg[offset : offset+int(length)]))
		domain.WriteString(".")
		offset += int(length)
	}
	host := domain.String()
	testHost := host
	parts := strings.Split(testHost, ".")
	try := 1
	for {
		if _, ok := blocked[testHost]; ok {
			block = true
			break
		}
		parts = parts[1:]
		if len(parts) < 3 {
			break
		}
		testHost = strings.Join(parts, ".")
		try++
	}

	if block {
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}
		cntBlocked.Add(1)

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		msg = append(msg, msg[12:12+1+len(host)]...) // domain
		msg = append(msg, answer...)                 // payload
		_, err := proxy.WriteTo(msg, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = &query{From: from, Host: host}
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			delete(queries, id)
			return
		}
		go func(queryID int) {
			time.Sleep(*flagTimeout)
			if query, ok := queries[queryID]; ok {
				fmt.Printf("DNS WARN: Query id %d %s timed out\n", queryID, query)
				cntTimedout.Add(1)
				delete(queries, queryID)
			}
			return
		}(id)
	}

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<|endoftext|>


<commit_before>package weatherupdate

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/urlfetch"
)

const (
	weddingYear  = 2017
	weddingMonth = time.September
	weddingDay   = 23
	weddingHour  = 15

	twilioNumber = "+18457648204"
	passphrase   = "wedding weather"

	weatherHistoryKey = "history1"
)

var (
	weatherAPI       = "http://api.openweathermap.org/data/2.5/forecast?zip=22304&units=imperial&APPID=" + os.Getenv("WEATHER_API_KEY")
	twilioSendSMSAPI = "https://api.twilio.com/2010-04-01/Accounts/" + os.Getenv("TWILIO_ACCOUNT_SID") + "/Messages"
)

func init() {
	http.HandleFunc("/checkTask", checkTaskHandler)
	http.HandleFunc("/checkNow", checkNowHandler)
	http.HandleFunc("/checkPhone", checkPhoneHandler)
}

type weatherResponse struct {
	List []struct {
		Dt   int64 `json:"dt"`
		Main struct {
			Temp     float32 `json:"temp"`
			Humidity int     `json:"humidity"`
		} `json:"main"`
		Weather []struct {
			ID          int    `json:"id"`
			Main        string `json:"main"`
			Description string `json:"description"`
			Icon        string `json:"icon"`
		} `json:"weather"`
	} `json:"list"`
}

func checkTaskHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	var history = new(weatherHistory)
	history.load(ctx)

	if history.changed(*results) {
		msgBody := messageBody(*results)
		history.update(ctx, *results)

		receiverNumbers := []string{"+18456496408", "+12405384163"}
		for _, n := range receiverNumbers {
			log.Infof(ctx, "sending to: %s", n)
			if err := sendMessage(ctx, msgBody, n); err != nil {
				http.Error(w, "could not send sms", http.StatusInternalServerError)
				log.Errorf(ctx, "could not send sms: %v", err)
				return
			}
		}
	}
}

const smsAllowContent = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Message>
			%s
		</Message>
	</Response>`

const smsDenyContent = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Message>
			<Media>%s</Media>
		</Message>
	</Response>`

func checkNowHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	w.Header().Set("Content-Type", "text/xml")

	body := r.PostFormValue("Body")
	if !strings.EqualFold(body, passphrase) {
		media := "http://www.icge.co.uk/languagesciencesblog/wp-content/uploads/2014/04/you_shall_not_pass1.jpg"
		fmt.Fprintf(w, smsDenyContent, media)
		return
	}

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	reply := messageBody(*results)
	fmt.Fprintf(w, smsAllowContent, reply)
}

const phoneMessage = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Say voice="woman">%s</Say>
	</Response>`

func checkPhoneHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	w.Header().Set("Content-Type", "text/xml")

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	reply := phoneBody(*results)
	fmt.Fprintf(w, phoneMessage, reply)
}

func fetchWeather(ctx context.Context, url string) (*weatherResponse, error) {
	resp, err := urlfetch.Client(ctx).Get(url)
	if err != nil {
		return nil, fmt.Errorf("could not get weather api data: %v", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf(resp.Status)
	}

	var results = new(weatherResponse)

	defer resp.Body.Close()
	if err = json.NewDecoder(resp.Body).Decode(&results); err != nil {
		return nil, fmt.Errorf("count not decode weather response: %v", err)
	}
	return results, nil
}

func sendMessage(ctx context.Context, body string, toNumber string) error {
	vals := url.Values{}
	vals.Set("To", toNumber)
	vals.Set("From", twilioNumber)
	vals.Set("Body", body)

	mb := strings.NewReader(vals.Encode())
	req, err := http.NewRequest("POST", twilioSendSMSAPI, mb)
	if err != nil {
		return fmt.Errorf("could not create sms request: %v", err)
	}
	req.SetBasicAuth(os.Getenv("TWILIO_ACCOUNT_SID"), os.Getenv("TWILIO_AUTH_TOKEN"))
	req.Header.Add("Accept", "application/json")
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	_, err = urlfetch.Client(ctx).Do(req)
	if err != nil {
		return fmt.Errorf("could not send sms request: %v", err)
	}
	return nil
}

func messageBody(results weatherResponse) string {
	msg := "Wedding Weather Report\n"
	for _, r := range results.List {
		t := time.Unix(r.Dt, 0)
		if isWeddingDay(t) {
			msg += fmt.Sprintf("Temp: %d degrees\nHumidity: %d%%\n%s: %s",
				int(r.Main.Temp), r.Main.Humidity, r.Weather[0].Main, r.Weather[0].Description)
		}
	}
	return msg
}

func phoneBody(results weatherResponse) string {
	msg := "Wedding Weather Report. September 23rd at 3pm. "
	for _, r := range results.List {
		t := time.Unix(r.Dt, 0)
		if isWeddingDay(t) {
			msg += fmt.Sprintf("Temperature: %d degrees. Humidity: %d%%. Conditions: %s.",
				int(r.Main.Temp), r.Main.Humidity, r.Weather[0].Description)
		}
	}
	return msg
}

func isWeddingDay(d time.Time) bool {
	return d.Year() == weddingYear && d.Month() == weddingMonth &&
		d.Day() == weddingDay && d.Hour() == weddingHour
}
<commit_msg>Updated code to not exit if failed to send message to one of the phone numbers.<commit_after>package weatherupdate

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/urlfetch"
)

const (
	weddingYear  = 2017
	weddingMonth = time.September
	weddingDay   = 23
	weddingHour  = 15

	twilioNumber = "+18457648204"
	passphrase   = "wedding weather"

	weatherHistoryKey = "history1"
)

var (
	weatherAPI       = "http://api.openweathermap.org/data/2.5/forecast?zip=22304&units=imperial&APPID=" + os.Getenv("WEATHER_API_KEY")
	twilioSendSMSAPI = "https://api.twilio.com/2010-04-01/Accounts/" + os.Getenv("TWILIO_ACCOUNT_SID") + "/Messages"
)

func init() {
	http.HandleFunc("/checkTask", checkTaskHandler)
	http.HandleFunc("/checkNow", checkNowHandler)
	http.HandleFunc("/checkPhone", checkPhoneHandler)
}

type weatherResponse struct {
	List []struct {
		Dt   int64 `json:"dt"`
		Main struct {
			Temp     float32 `json:"temp"`
			Humidity int     `json:"humidity"`
		} `json:"main"`
		Weather []struct {
			ID          int    `json:"id"`
			Main        string `json:"main"`
			Description string `json:"description"`
			Icon        string `json:"icon"`
		} `json:"weather"`
	} `json:"list"`
}

func checkTaskHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	var history = new(weatherHistory)
	history.load(ctx)

	if history.changed(*results) {
		msgBody := messageBody(*results)
		history.update(ctx, *results)

		receiverNumbers := []string{"+18456496408", "+12405384163"}
		for _, n := range receiverNumbers {
			log.Infof(ctx, "sending to: %s", n)
			if err := sendMessage(ctx, msgBody, n); err != nil {
				http.Error(w, "could not send sms", http.StatusInternalServerError)
				log.Errorf(ctx, "could not send sms: %v", err)
			}
		}
	}
}

const smsAllowContent = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Message>
			%s
		</Message>
	</Response>`

const smsDenyContent = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Message>
			<Media>%s</Media>
		</Message>
	</Response>`

func checkNowHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	w.Header().Set("Content-Type", "text/xml")

	body := r.PostFormValue("Body")
	if !strings.EqualFold(body, passphrase) {
		media := "http://www.icge.co.uk/languagesciencesblog/wp-content/uploads/2014/04/you_shall_not_pass1.jpg"
		fmt.Fprintf(w, smsDenyContent, media)
		return
	}

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	reply := messageBody(*results)
	fmt.Fprintf(w, smsAllowContent, reply)
}

const phoneMessage = `<?xml version="1.0" encoding="UTF-8"?>
	<Response>
		<Say voice="woman">%s</Say>
	</Response>`

func checkPhoneHandler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	w.Header().Set("Content-Type", "text/xml")

	results, err := fetchWeather(ctx, weatherAPI)
	if err != nil {
		http.Error(w, "could not fetch weather data", http.StatusInternalServerError)
		log.Errorf(ctx, "could not fetch weather data: %v", err)
		return
	}

	reply := phoneBody(*results)
	fmt.Fprintf(w, phoneMessage, reply)
}

func fetchWeather(ctx context.Context, url string) (*weatherResponse, error) {
	resp, err := urlfetch.Client(ctx).Get(url)
	if err != nil {
		return nil, fmt.Errorf("could not get weather api data: %v", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf(resp.Status)
	}

	var results = new(weatherResponse)

	defer resp.Body.Close()
	if err = json.NewDecoder(resp.Body).Decode(&results); err != nil {
		return nil, fmt.Errorf("count not decode weather response: %v", err)
	}
	return results, nil
}

func sendMessage(ctx context.Context, body string, toNumber string) error {
	vals := url.Values{}
	vals.Set("To", toNumber)
	vals.Set("From", twilioNumber)
	vals.Set("Body", body)

	mb := strings.NewReader(vals.Encode())
	req, err := http.NewRequest("POST", twilioSendSMSAPI, mb)
	if err != nil {
		return fmt.Errorf("could not create sms request: %v", err)
	}
	req.SetBasicAuth(os.Getenv("TWILIO_ACCOUNT_SID"), os.Getenv("TWILIO_AUTH_TOKEN"))
	req.Header.Add("Accept", "application/json")
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	_, err = urlfetch.Client(ctx).Do(req)
	if err != nil {
		return fmt.Errorf("could not send sms request: %v", err)
	}
	return nil
}

func messageBody(results weatherResponse) string {
	msg := "Wedding Weather Report\n"
	for _, r := range results.List {
		t := time.Unix(r.Dt, 0)
		if isWeddingDay(t) {
			msg += fmt.Sprintf("Temp: %d degrees\nHumidity: %d%%\n%s: %s",
				int(r.Main.Temp), r.Main.Humidity, r.Weather[0].Main, r.Weather[0].Description)
		}
	}
	return msg
}

func phoneBody(results weatherResponse) string {
	msg := "Wedding Weather Report. September 23rd at 3pm. "
	for _, r := range results.List {
		t := time.Unix(r.Dt, 0)
		if isWeddingDay(t) {
			msg += fmt.Sprintf("Temperature: %d degrees. Humidity: %d%%. Conditions: %s.",
				int(r.Main.Temp), r.Main.Humidity, r.Weather[0].Description)
		}
	}
	return msg
}

func isWeddingDay(d time.Time) bool {
	return d.Year() == weddingYear && d.Month() == weddingMonth &&
		d.Day() == weddingDay && d.Hour() == weddingHour
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"sync"

	"strconv"

	_ "github.com/heroku/x/hmetrics"
)

type fileDef struct {
	path string
	url  string
}
type remoteFile struct {
	path     string
	response *http.Response
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "7777"
	}
	log.Println("Listening on port " + port)
	http.ListenAndServe(":"+port, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		fileName := r.FormValue("filename")
		if fileName == "" {
			fileName = "download.zip"
		}
		w.Header().Add("Content-Disposition", "attachment; filename=\""+fileName+"\"")
		log.Println("Generating zip for " + fileName)

		archive := zip.NewWriter(w)

		downloadQueue := make(chan fileDef)
		archiveQueue := make(chan remoteFile)
		downloaderWaitGroup := sync.WaitGroup{}
		finished := make(chan error)

		// Run archiver
		go func() {
			for rf := range archiveQueue {
				log.Println("Received data for " + rf.path)
				entryHeader := &zip.FileHeader{
					Name:   rf.path,
					Method: zip.Deflate,
				}
				entryHeader.SetModTime(time.Now())
				entryHeader.SetMode(os.ModePerm)
				entry, err := archive.CreateHeader(entryHeader)
				if err != nil {
					finished <- err
					return
				}
				log.Println(rf.response.Header.Get("Last-Modified"))
				modifiedTime, err := time.Parse("Mon, 2 Jan 2006 15:04:05 MST", rf.response.Header.Get("Last-Modified"))
				if err != nil {
					entryHeader.SetModTime(modifiedTime)
				}

				io.Copy(entry, rf.response.Body)
				rf.response.Body.Close()
				log.Println("Zipped and closed " + rf.path)
			}
			log.Println("All zipping complete, closing up.")
			err := archive.Close()
			if err != nil {
				finished <- err
				return
			}
			log.Println("Acrhive closed.")
			finished <- nil
		}()

		// Run downloaders
		for w := 1; w <= 5; w++ {
			go func(w int) {
				log.Println("Downloader Ready: " + strconv.Itoa(w))
				for df := range downloadQueue {
					log.Println("Downloading[" + strconv.Itoa(w) + "] path: " + df.path + " URL: " + df.url)
					download, err := http.Get(df.url)
					if err != nil {
						downloaderWaitGroup.Done()
						finished <- err
						return
					}
					archiveQueue <- remoteFile{path: df.path, response: download}
					downloaderWaitGroup.Done()
				}
				log.Println("Downloader Shutting Down: " + strconv.Itoa(w))
			}(w)
		}

		for path, urls := range r.PostForm {
			downloaderWaitGroup.Add(1)
			downloadQueue <- fileDef{path: path, url: urls[0]}
		}
		close(downloadQueue)
		downloaderWaitGroup.Wait()
		close(archiveQueue)
		err := <-finished
		if err != nil {
			handleError(err, w)
			return
		}
	}))
}

func handleError(err error, w http.ResponseWriter) {
	http.Error(w, "error", http.StatusBadRequest)
	log.Println(err)
	return
}
<commit_msg>Set env var for concurrency<commit_after>package main

import (
	"archive/zip"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"sync"

	"strconv"

	_ "github.com/heroku/x/hmetrics"
)

type fileDef struct {
	path string
	url  string
}
type remoteFile struct {
	path     string
	response *http.Response
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "7777"
	}
	concurrency, err := strconv.Atoi(os.Getenv("DOWNLOAD_CONCURRENCY"))
	if err != nil {
		concurrency = 10
	}

	log.Println("Listening on port " + port)
	http.ListenAndServe(":"+port, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		fileName := r.FormValue("filename")
		if fileName == "" {
			fileName = "download.zip"
		}
		w.Header().Add("Content-Disposition", "attachment; filename=\""+fileName+"\"")
		log.Println("Generating zip for " + fileName)

		archive := zip.NewWriter(w)

		downloadQueue := make(chan fileDef)
		archiveQueue := make(chan remoteFile)
		downloaderWaitGroup := sync.WaitGroup{}
		finished := make(chan error)

		// Run archiver
		go func() {
			for rf := range archiveQueue {
				log.Println("Received data for " + rf.path)
				entryHeader := &zip.FileHeader{
					Name:   rf.path,
					Method: zip.Deflate,
				}
				entryHeader.SetModTime(time.Now())
				entryHeader.SetMode(os.ModePerm)
				entry, err := archive.CreateHeader(entryHeader)
				if err != nil {
					finished <- err
					return
				}
				log.Println(rf.response.Header.Get("Last-Modified"))
				modifiedTime, err := time.Parse("Mon, 2 Jan 2006 15:04:05 MST", rf.response.Header.Get("Last-Modified"))
				if err != nil {
					entryHeader.SetModTime(modifiedTime)
				}

				io.Copy(entry, rf.response.Body)
				rf.response.Body.Close()
				log.Println("Zipped and closed " + rf.path)
			}
			log.Println("All zipping complete, closing up.")
			err := archive.Close()
			if err != nil {
				finished <- err
				return
			}
			log.Println("Acrhive closed.")
			finished <- nil
		}()

		// Run downloaders
		for w := 1; w <= concurrency; w++ {
			go func(w int) {
				log.Println("Downloader Ready: " + strconv.Itoa(w))
				for df := range downloadQueue {
					log.Println("Downloading[" + strconv.Itoa(w) + "] path: " + df.path + " URL: " + df.url)
					download, err := http.Get(df.url)
					if err != nil {
						downloaderWaitGroup.Done()
						finished <- err
						return
					}
					archiveQueue <- remoteFile{path: df.path, response: download}
					downloaderWaitGroup.Done()
				}
				log.Println("Downloader Shutting Down: " + strconv.Itoa(w))
			}(w)
		}

		for path, urls := range r.PostForm {
			downloaderWaitGroup.Add(1)
			downloadQueue <- fileDef{path: path, url: urls[0]}
		}
		close(downloadQueue)
		downloaderWaitGroup.Wait()
		close(archiveQueue)
		err := <-finished
		if err != nil {
			handleError(err, w)
			return
		}
	}))
}

func handleError(err error, w http.ResponseWriter) {
	http.Error(w, "error", http.StatusBadRequest)
	log.Println(err)
	return
}
<|endoftext|>


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

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/venicegeo/pzsvc-exec/pzsvc"
)

type configType struct {
	CliCmd      string
	PzAddr      string
	SvcName     string
	URL         string
	Port        int
	Description	string
	ImageReqs	map[string]string
	Attributes	map[string]string
}

type outStruct struct {
	InFiles		map[string]string
	OutFiles	map[string]string
	ProgReturn	string
	Errors		[]string
}

func main() {

	if len(os.Args) < 2 {
		fmt.Println("error: Insufficient parameters.  You must specify a config file.")
		return
	}

	// first argument after the base call should be the path to the config file.
	// ReadFile returns the contents of the file as a byte buffer.
	configBuf, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Println("error:", err)
	}

	var configObj configType
	err = json.Unmarshal(configBuf, &configObj)
	if err != nil {
		fmt.Println("error:", err)
	}

	if configObj.Port <= 0 {
		configObj.Port = 8080
	}

	portStr := ":" + strconv.Itoa(configObj.Port)

	if configObj.SvcName != "" && configObj.PzAddr != "" {
		fmt.Println("About to manage registration.")
		err = pzsvc.ManageRegistration(	configObj.SvcName,
										configObj.Description,
										configObj.URL + "/execute",
										configObj.PzAddr,
										configObj.ImageReqs )
		if err != nil {
			fmt.Println("error:", err.Error())
		}
		fmt.Println("Registration managed.")

	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		switch r.URL.Path {
		case "/":
			fmt.Fprintf(w, "hello.")
		case "/execute":
			{

				var cmdParam string
				var inFileStr string
				var outTiffStr string
				var outTxtStr string
				var outGeoJStr string
				var usePz string

				// might be time to start looking into that "help" thing.

				if r.Method == "GET" {
					cmdParam = r.URL.Query().Get("cmd")
					inFileStr = r.URL.Query().Get("inFiles")
					outTiffStr = r.URL.Query().Get("outTiffs")
					outTxtStr = r.URL.Query().Get("outTxts")
					outGeoJStr = r.URL.Query().Get("outGeoJson")
					usePz = r.URL.Query().Get("pz")
				} else {
					cmdParam = r.FormValue("cmd")
					inFileStr = r.FormValue("inFiles")
					outTiffStr = r.FormValue("outTiffs")
					outTxtStr = r.FormValue("outTxts")
					outGeoJStr = r.FormValue("outGeoJson")
					usePz = r.FormValue("pz")
				}

				cmdConfigSlice := splitOrNil(configObj.CliCmd, " ")
				cmdParamSlice := splitOrNil(cmdParam, " ")
				cmdSlice := append(cmdConfigSlice, cmdParamSlice...)

				inFileSlice := splitOrNil(inFileStr, ",")
				outTiffSlice := splitOrNil(outTiffStr, ",")
				outTxtSlice := splitOrNil(outTxtStr, ",")
				outGeoJSlice := splitOrNil(outGeoJStr, ",")

				var output outStruct

				runID, err := psuUUID()
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}

				err = os.Mkdir("./"+runID, 0777)
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}
				defer os.RemoveAll("./" + runID)

				err = os.Chmod("./"+runID, 0777)
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}

				if len(inFileSlice) > 0 {
					output.InFiles = make(map[string]string)
				}
				if len(outTiffSlice)+len(outTxtSlice)+len(outGeoJSlice) > 0 {
					output.OutFiles = make(map[string]string)
				}

				for i, inFile := range inFileSlice {

					fmt.Printf("Downloading file %s - %d of %d.\n", inFile, i, len(inFileSlice))
					fname, err := pzsvc.Download(inFile, runID, configObj.PzAddr)
					if err != nil {
						output.Errors = append(output.Errors,err.Error())
						fmt.Printf("Download failed.  %s", err.Error())
					} else {
						output.InFiles[inFile] = fname
						fmt.Printf("Successfully downloaded %s.", fname)
					}
				}

				if len(cmdSlice) == 0 {
					output.Errors = append(output.Errors, `No cmd specified in config file.  Please provide "cmd" param.`)
					break
				}

				fmt.Printf("Executing \"%s\".\n", configObj.CliCmd+" "+cmdParam)

				// we're calling this from inside a temporary subfolder.  If the
				// program called exists inside the initial pzsvc-exec folder, that's
				// probably where it's called from, and we need to acccess it directly.
				_, err = os.Stat(fmt.Sprintf("./%s", cmdSlice[0]))
				if err == nil || !(os.IsNotExist(err)){
					// ie, if there's a file in the start folder named the same thing
					// as the base command
					cmdSlice[0] = ("../" + cmdSlice[0])
				}

				clc := exec.Command(cmdSlice[0], cmdSlice[1:]...)
				clc.Dir = runID

				var b bytes.Buffer
				clc.Stdout = &b
				clc.Stderr = os.Stderr

				err = clc.Run()
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}
				fmt.Printf("Program output: %+q\n", b.String())
				output.ProgReturn = fmt.Sprintf("%+q", b.String())

				for i, outTiff := range outTiffSlice {
					fmt.Printf("Uploading Tiff %s - %d of %d.\n", outTiff, i, len(outTiffSlice))
					dataID, err := pzsvc.IngestTiff(outTiff, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						fmt.Fprintf(w, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outTiff] = dataID
					}
				}

				for i, outTxt := range outTxtSlice {
					fmt.Printf("Uploading Txt %s - %d of %d.\n", outTxt, i, len(outTxtSlice))
					dataID, err := pzsvc.IngestTxt(outTxt, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outTxt] = dataID
					}
				}

				for i, outGeoJ := range outGeoJSlice {
					fmt.Printf("Uploading GeoJson %s - %d of %d.\n", outGeoJ, i, len(outGeoJSlice))
					dataID, err := pzsvc.IngestGeoJson(outGeoJ, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outGeoJ] = dataID
					}
				}

				// this isn't a thing now, but might again become a thing later.
				if usePz != "" {

					type pzCont struct {
						Type		string
						Content		string
						MimeType	string
					}

					type pzWrap struct {
						DataType	pzCont
						metadata	map[string]string
					}

					var cont pzCont
					var wrap pzWrap
					outBuf, err := json.Marshal(output)
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
					}

					cont.Type = "text"
					cont.Content = string(outBuf)
					cont.MimeType = "text/plain"

					wrap.DataType = cont

					printJson(w, wrap)
				} else {
					printJson(w, output)
				}

			}
		case "/description":
			if configObj.Description == "" {
				fmt.Fprintf(w, "No description defined")
			} else {
				fmt.Fprintf(w, configObj.Description)
			}
		case "/attributes":
			if configObj.Attributes == nil {
				fmt.Fprintf(w, "{ }")
			} else {
				printJson(w, configObj.Attributes)
			}
		case "/help":
			fmt.Fprintf(w, "We're sorry, help is not yet implemented.\n")
		default:
			fmt.Fprintf(w, "Command undefined.  Try help?\n")
		}
	})

	log.Fatal(http.ListenAndServe(portStr, nil))
}

func splitOrNil(inString, knife string) []string {
	if inString == "" {
		return nil
	}
	return strings.Split(inString, knife)
}

func psuUUID() (string, error) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]), nil
}

func printJson (w http.ResponseWriter, output interface{}) {
	outBuf, err := json.Marshal(output)
	if err != nil {
		fmt.Fprintf(w, `{"Errors":"Json marshalling failure.  Data not reportable."}`)
	}

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

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/venicegeo/pzsvc-exec/pzsvc"
)

type configType struct {
	CliCmd      string
	PzAddr      string
	SvcName     string
	URL         string
	Port        int
	Description	string
	ImageReqs	map[string]string
	Attributes	map[string]string
}

type outStruct struct {
	InFiles		map[string]string
	OutFiles	map[string]string
	ProgReturn	string
	Errors		[]string
}

func main() {

	if len(os.Args) < 2 {
		fmt.Println("error: Insufficient parameters.  You must specify a config file.")
		return
	}

	// first argument after the base call should be the path to the config file.
	// ReadFile returns the contents of the file as a byte buffer.
	configBuf, err := ioutil.ReadFile(os.Args[1])
	if err != nil {
		fmt.Println("error:", err)
	}

	var configObj configType
	err = json.Unmarshal(configBuf, &configObj)
	if err != nil {
		fmt.Println("error:", err)
	}

	if configObj.Port <= 0 {
		configObj.Port = 8080
	}

	portStr := ":" + strconv.Itoa(configObj.Port)

	if configObj.SvcName != "" && configObj.PzAddr != "" {
		fmt.Println("About to manage registration.")
		err = pzsvc.ManageRegistration(	configObj.SvcName,
										configObj.Description,
										configObj.URL + "/execute",
										configObj.PzAddr,
										configObj.ImageReqs )
		if err != nil {
			fmt.Println("error:", err.Error())
		}
		fmt.Println("Registration managed.")

	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		switch r.URL.Path {
		case "/":
			fmt.Fprintf(w, "hello.")
		case "/execute":
			{

				var cmdParam string
				var inFileStr string
				var outTiffStr string
				var outTxtStr string
				var outGeoJStr string
				var usePz string

				// might be time to start looking into that "help" thing.

				if r.Method == "GET" {
					cmdParam = r.URL.Query().Get("cmd")
					inFileStr = r.URL.Query().Get("inFiles")
					outTiffStr = r.URL.Query().Get("outTiffs")
					outTxtStr = r.URL.Query().Get("outTxts")
					outGeoJStr = r.URL.Query().Get("outGeoJson")
					usePz = r.URL.Query().Get("pz")
				} else {
					cmdParam = r.FormValue("cmd")
					inFileStr = r.FormValue("inFiles")
					outTiffStr = r.FormValue("outTiffs")
					outTxtStr = r.FormValue("outTxts")
					outGeoJStr = r.FormValue("outGeoJson")
					usePz = r.FormValue("pz")
				}

				cmdConfigSlice := splitOrNil(configObj.CliCmd, " ")
				cmdParamSlice := splitOrNil(cmdParam, " ")
				cmdSlice := append(cmdConfigSlice, cmdParamSlice...)

				inFileSlice := splitOrNil(inFileStr, ",")
				outTiffSlice := splitOrNil(outTiffStr, ",")
				outTxtSlice := splitOrNil(outTxtStr, ",")
				outGeoJSlice := splitOrNil(outGeoJStr, ",")

				var output outStruct

				runID, err := psuUUID()
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}

				err = os.Mkdir("./"+runID, 0777)
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}
				defer os.RemoveAll("./" + runID)

				err = os.Chmod("./"+runID, 0777)
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}

				if len(inFileSlice) > 0 {
					output.InFiles = make(map[string]string)
				}
				if len(outTiffSlice)+len(outTxtSlice)+len(outGeoJSlice) > 0 {
					output.OutFiles = make(map[string]string)
				}

				for i, inFile := range inFileSlice {

					fmt.Printf("Downloading file %s - %d of %d.\n", inFile, i, len(inFileSlice))
					fname, err := pzsvc.Download(inFile, runID, configObj.PzAddr)
					if err != nil {
						output.Errors = append(output.Errors,err.Error())
						fmt.Printf("Download failed.  %s", err.Error())
					} else {
						output.InFiles[inFile] = fname
						fmt.Printf("Successfully downloaded %s.", fname)
					}
				}

				if len(cmdSlice) == 0 {
					output.Errors = append(output.Errors, `No cmd specified in config file.  Please provide "cmd" param.`)
					break
				}

				fmt.Printf("Executing \"%s\".\n", configObj.CliCmd+" "+cmdParam)

				// we're calling this from inside a temporary subfolder.  If the
				// program called exists inside the initial pzsvc-exec folder, that's
				// probably where it's called from, and we need to acccess it directly.
				_, err = os.Stat(fmt.Sprintf("./%s", cmdSlice[0]))
				if err == nil || !(os.IsNotExist(err)){
					// ie, if there's a file in the start folder named the same thing
					// as the base command
					cmdSlice[0] = ("../" + cmdSlice[0])
				}

				clc := exec.Command(cmdSlice[0], cmdSlice[1:]...)
				clc.Dir = runID

				var b bytes.Buffer
				clc.Stdout = &b
				clc.Stderr = os.Stderr

				err = clc.Run()
				if err != nil {
					output.Errors = append(output.Errors, err.Error())
				}
				output.ProgReturn = b.String())
				output.ProgReturn = strings.Replace(output.ProgReturn, "\r\n", "\n", -1)
				output.ProgReturn = strings.Replace(output.ProgReturn, "\r", "\n", -1)

				fmt.Printf("Program output: %s\n", output.ProgReturn)

				for i, outTiff := range outTiffSlice {
					fmt.Printf("Uploading Tiff %s - %d of %d.\n", outTiff, i, len(outTiffSlice))
					dataID, err := pzsvc.IngestTiff(outTiff, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						fmt.Fprintf(w, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outTiff] = dataID
					}
				}

				for i, outTxt := range outTxtSlice {
					fmt.Printf("Uploading Txt %s - %d of %d.\n", outTxt, i, len(outTxtSlice))
					dataID, err := pzsvc.IngestTxt(outTxt, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outTxt] = dataID
					}
				}

				for i, outGeoJ := range outGeoJSlice {
					fmt.Printf("Uploading GeoJson %s - %d of %d.\n", outGeoJ, i, len(outGeoJSlice))
					dataID, err := pzsvc.IngestGeoJson(outGeoJ, runID, configObj.PzAddr, cmdSlice[0])
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
						fmt.Printf("Upload failed.  %s", err.Error())
					} else {
						output.OutFiles[outGeoJ] = dataID
					}
				}

				// this isn't a thing now, but might again become a thing later.
				if usePz != "" {

					type pzCont struct {
						Type		string
						Content		string
						MimeType	string
					}

					type pzWrap struct {
						DataType	pzCont
						metadata	map[string]string
					}

					var cont pzCont
					var wrap pzWrap
					outBuf, err := json.Marshal(output)
					if err != nil {
						output.Errors = append(output.Errors, err.Error())
					}

					cont.Type = "text"
					cont.Content = string(outBuf)
					cont.MimeType = "text/plain"

					wrap.DataType = cont

					printJson(w, wrap)
				} else {
					printJson(w, output)
				}

			}
		case "/description":
			if configObj.Description == "" {
				fmt.Fprintf(w, "No description defined")
			} else {
				fmt.Fprintf(w, configObj.Description)
			}
		case "/attributes":
			if configObj.Attributes == nil {
				fmt.Fprintf(w, "{ }")
			} else {
				printJson(w, configObj.Attributes)
			}
		case "/help":
			fmt.Fprintf(w, "We're sorry, help is not yet implemented.\n")
		default:
			fmt.Fprintf(w, "Command undefined.  Try help?\n")
		}
	})

	log.Fatal(http.ListenAndServe(portStr, nil))
}

func splitOrNil(inString, knife string) []string {
	if inString == "" {
		return nil
	}
	return strings.Split(inString, knife)
}

func psuUUID() (string, error) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]), nil
}

func printJson (w http.ResponseWriter, output interface{}) {
	outBuf, err := json.Marshal(output)
	if err != nil {
		fmt.Fprintf(w, `{"Errors":"Json marshalling failure.  Data not reportable."}`)
	}

	outStr := string(outBuf)
	fmt.Fprintf(w, outStr)
}<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"

	"github.com/MyHomeworkSpace/api-server/schools"
	"github.com/julienschmidt/httprouter"

	"github.com/MyHomeworkSpace/api-server/schools/cornell"
	"github.com/MyHomeworkSpace/api-server/schools/dalton"
	"github.com/MyHomeworkSpace/api-server/schools/mit"

	"github.com/MyHomeworkSpace/api-server/api"
	"github.com/MyHomeworkSpace/api-server/auth"
	"github.com/MyHomeworkSpace/api-server/config"
	"github.com/MyHomeworkSpace/api-server/data"
	"github.com/MyHomeworkSpace/api-server/email"
)

type errorResponse struct {
	Status string `json:"status"`
	Error  string `json:"error"`
}

type csrfResponse struct {
	Status string `json:"status"`
	Token  string `json:"token"`
}

func main() {
	log.Println("MyHomeworkSpace API Server")

	config.Init()

	initDatabase()
	initRedis()
	initWebauthn()

	migrationName := flag.String("migrate", "", "If specified, the API server will run the migration with the given name.")
	flag.Parse()

	if *migrationName != "" {
		migrate(*migrationName)
		return
	}

	email.Init()

	api.DB = DB
	api.MainRegistry = schools.MainRegistry
	api.RedisClient = RedisClient
	api.WebAuthnHandler = WebAuthnHandler

	auth.DB = DB
	auth.RedisClient = RedisClient

	data.DB = DB
	data.MainRegistry = schools.MainRegistry
	data.RedisClient = RedisClient

	schools.MainRegistry.Register(dalton.CreateSchool())
	schools.MainRegistry.Register(mit.CreateSchool())
	schools.MainRegistry.Register(cornell.CreateSchool())

	router := httprouter.New()

	router.GET("/", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		w.Write([]byte("MyHomeworkSpace API Server"))
	})
	router.ServeFiles("/api_tester/*filepath", http.Dir("api_tester/"))

	api.Init(router) // API init delayed because router must be started first
	http.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// if it's a preflight, handle it here
		if r.Method == "OPTIONS" {
			w.Header().Set("Access-Control-Allow-Credentials", "false")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Access-Control-Allow-Headers", "authorization")
			w.WriteHeader(http.StatusOK)
			return
		}

		// otherwise, pass it through
		router.ServeHTTP(w, r)
	}))

	log.Printf("Listening on port %d", config.GetCurrent().Server.Port)
	err := http.ListenAndServe(fmt.Sprintf(":%d", config.GetCurrent().Server.Port), nil)
	if err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>temporarily disable webauthn code<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"

	"github.com/MyHomeworkSpace/api-server/schools"
	"github.com/julienschmidt/httprouter"

	"github.com/MyHomeworkSpace/api-server/schools/cornell"
	"github.com/MyHomeworkSpace/api-server/schools/dalton"
	"github.com/MyHomeworkSpace/api-server/schools/mit"

	"github.com/MyHomeworkSpace/api-server/api"
	"github.com/MyHomeworkSpace/api-server/auth"
	"github.com/MyHomeworkSpace/api-server/config"
	"github.com/MyHomeworkSpace/api-server/data"
	"github.com/MyHomeworkSpace/api-server/email"
)

type errorResponse struct {
	Status string `json:"status"`
	Error  string `json:"error"`
}

type csrfResponse struct {
	Status string `json:"status"`
	Token  string `json:"token"`
}

func main() {
	log.Println("MyHomeworkSpace API Server")

	config.Init()

	initDatabase()
	initRedis()
	// initWebauthn()

	migrationName := flag.String("migrate", "", "If specified, the API server will run the migration with the given name.")
	flag.Parse()

	if *migrationName != "" {
		migrate(*migrationName)
		return
	}

	email.Init()

	api.DB = DB
	api.MainRegistry = schools.MainRegistry
	api.RedisClient = RedisClient
	api.WebAuthnHandler = WebAuthnHandler

	auth.DB = DB
	auth.RedisClient = RedisClient

	data.DB = DB
	data.MainRegistry = schools.MainRegistry
	data.RedisClient = RedisClient

	schools.MainRegistry.Register(dalton.CreateSchool())
	schools.MainRegistry.Register(mit.CreateSchool())
	schools.MainRegistry.Register(cornell.CreateSchool())

	router := httprouter.New()

	router.GET("/", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		w.Write([]byte("MyHomeworkSpace API Server"))
	})
	router.ServeFiles("/api_tester/*filepath", http.Dir("api_tester/"))

	api.Init(router) // API init delayed because router must be started first
	http.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// if it's a preflight, handle it here
		if r.Method == "OPTIONS" {
			w.Header().Set("Access-Control-Allow-Credentials", "false")
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Access-Control-Allow-Headers", "authorization")
			w.WriteHeader(http.StatusOK)
			return
		}

		// otherwise, pass it through
		router.ServeHTTP(w, r)
	}))

	log.Printf("Listening on port %d", config.GetCurrent().Server.Port)
	err := http.ListenAndServe(fmt.Sprintf(":%d", config.GetCurrent().Server.Port), nil)
	if err != nil {
		log.Fatalln(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/gin-gonic/gin"
)

type Host struct {
	Hostname   string    `json:"hostname"`
	Expiration time.Time `json:"expiration"`
}

var (
	hosts []Host
	TTL   = 30
)

func buildExpiration() time.Time {
	expiration := time.Now().UTC()
	expiration = expiration.Add(time.Duration(TTL) * time.Second)

	return expiration
}

func cleanup() {
	for {
		for i, h := range hosts {
			if time.Now().UTC().After(h.Expiration) {
				diff := hosts
				diff = append(diff[:i], diff[i+1:]...)
				hosts = diff
			}
		}
		time.Sleep(time.Duration(TTL) * time.Second)
	}
}

func main() {
	go cleanup()

	router := gin.Default()

	router.POST("/checkin/", func(c *gin.Context) {
		var host Host

		c.Bind(&host)

		if host.Hostname != "" {
			exists := false

			for i, h := range hosts {
				if h.Hostname == host.Hostname {
					exists = true
					hosts[i].Expiration = buildExpiration()
				}
			}

			if !exists {
				host.Expiration = buildExpiration()
				hosts = append(hosts, host)
			}

			c.JSON(200, "")
		} else {
			var response struct {
				Message string `json:"message"`
			}
			response.Message = "Invalid hostname."
			c.JSON(400, response)
		}
	})

	router.GET("/status/", func(c *gin.Context) {
		c.JSON(200, hosts)
	})

	router.Run(":8080")
}
<commit_msg>Added a Status struct to represent the current status<commit_after>package main

import (
	"time"

	"github.com/gin-gonic/gin"
)

type Host struct {
	Hostname   string    `json:"hostname"`
	Expiration time.Time `json:"expiration"`
}

type Status struct {
	Hosts []Host `json:"hosts"`
}

var (
	hosts []Host
	TTL   = 30
)

func buildExpiration() time.Time {
	expiration := time.Now().UTC()
	expiration = expiration.Add(time.Duration(TTL) * time.Second)

	return expiration
}

func cleanup() {
	for {
		for i, h := range hosts {
			if time.Now().UTC().After(h.Expiration) {
				diff := hosts
				diff = append(diff[:i], diff[i+1:]...)
				hosts = diff
			}
		}
		time.Sleep(time.Duration(TTL) * time.Second)
	}
}

func main() {
	go cleanup()

	router := gin.Default()

	router.POST("/checkin/", func(c *gin.Context) {
		var host Host

		c.Bind(&host)

		if host.Hostname != "" {
			exists := false

			for i, h := range hosts {
				if h.Hostname == host.Hostname {
					exists = true
					hosts[i].Expiration = buildExpiration()
				}
			}

			if !exists {
				host.Expiration = buildExpiration()
				hosts = append(hosts, host)
			}

			c.JSON(200, "")
		} else {
			var response struct {
				Message string `json:"message"`
			}
			response.Message = "Invalid hostname."
			c.JSON(400, response)
		}
	})

	router.GET("/status/", func(c *gin.Context) {
		status := Status{}
		status.Hosts = hosts

		c.JSON(200, status)
	})

	router.Run(":8080")
}
<|endoftext|>


<commit_before>// Creates a webserver/unix socket that allows a Docker server to create RDMA
// backed volumes.
package main

import (
	"errors"
	"flag"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/docker/go-plugins-helpers/volume"
	"github.com/golang/glog"
	"github.com/mellanox-senior-design/docker-volume-rdma/db"
	"github.com/mellanox-senior-design/docker-volume-rdma/drivers"
)

// Port to launch service on.
var httpPort int

// Volume Flags
var volumeDatabaseDriver string
var volumeDatabasePath string
var volumeDatabaseHost string
var volumeDatabaseUsername string
var volumeDatabasePassword string
var volumeDatabaseSchema string

// Storage Flags
var storageControllerDriver string
var storageControllerPath string

func init() {
	// Configure application flags.
	flag.IntVar(&httpPort, "port", 8080, "tcp/ip port to serve volume driver on")

	// Volume Database Flags
	flag.StringVar(&volumeDatabaseDriver, "db", "sqlite", "set the database backend used to store volume metadata: [sqlite, mysql, in-memory]")
	flag.StringVar(&volumeDatabasePath, "dbpath", "", "set the database storage path")
	flag.StringVar(&volumeDatabaseHost, "dbhost", "", "set the database host (default is '' or localhost:3306)")
	flag.StringVar(&volumeDatabaseUsername, "dbuser", "", "set the database username (default is root)")
	flag.StringVar(&volumeDatabasePassword, "dbpass", "", "set the database password (optional)")
	flag.StringVar(&volumeDatabaseSchema, "dbschema", "", "set the database schema (required)")

	// Storage Controller Flags
	flag.StringVar(&storageControllerDriver, "sc", "glusterfs", "set the storage backend used to store volume data: [glusterfs, on-disk]")
	flag.StringVar(&storageControllerPath, "scpath", "", "set the storage path used to know where to put the volumes on the host")
}

// Configure and start the docker volume plugin server.
func main() {

	// Parse flags as glog needs the flags to be solidified before starting.
	flag.Parse()

	// Convert port to string, and print startup message.
	port := strconv.Itoa(httpPort)

	driver, handler, err := configure()
	if err == nil {

		// Handle SIGINT gracefully
		sigintChannel := make(chan os.Signal, 2)
		signal.Notify(sigintChannel, os.Interrupt, syscall.SIGINT)
		go func() {
			<-sigintChannel
			glog.Info("Exiting ...")
			driver.Disconnect()
			os.Exit(1)
		}()

		err = driver.Connect()
		if err == nil {
			defer driver.Disconnect()

			glog.Info("Running! http://localhost:" + port)
			err = handler.ServeTCP("docker-volume-rdma", ":"+port, nil)
		}
	}

	// Log any error that may have occurred.
	glog.Fatal(err)
}

func configure() (*drivers.RDMAVolumeDriver, *volume.Handler, error) {
	// Create and begin serving volume driver on tcp/ip port, httpPort.
	volumeDatabase, err := getDatabaseConnection()
	if err != nil {
		return nil, nil, err
	}

	// Configure Storage Controller
	storageController, err := getStorageConnection()
	if err != nil {
		return nil, nil, err
	}

	// Print startup message and start server
	glog.Info("Connecting to services ...")
	driver := drivers.NewRDMAVolumeDriver(storageController, volumeDatabase)
	handler := volume.NewHandler(driver)

	return &driver, handler, nil
}

// GetDatabaseConnection returns the database connection that was requested on the command line.
func getDatabaseConnection() (db.VolumeDatabase, error) {
	glog.Info("Attempting to use the ", volumeDatabaseDriver, " volume driver.")
	switch volumeDatabaseDriver {
	case "in-memory":
		validateDatabaseFlags(true, true, false, false, false, false)
		return db.NewInMemoryVolumeDatabase(), nil

	case "sqlite":
		validateDatabaseFlags(false, true, false, false, false, false)
		return db.NewSQLiteVolumeDatabase(volumeDatabasePath), nil

	case "mysql":
		validateDatabaseFlags(false, false, true, true, true, true)
		return db.NewMySQLVolumeDatabase(volumeDatabaseHost, volumeDatabaseUsername, volumeDatabasePassword, volumeDatabaseSchema)

	default:
		return nil, errors.New("unsupported database, please choose sqlite, mysql, or in-memory")
	}
}

func validateDatabaseFlags(fatal bool, path bool, host bool, username bool, password bool, schema bool) {

	var errors bool
	noteErrorFunc := func(name string, value string, used bool) {
		if !used && value != "" {
			glog.Warning("Volume Driver: ", volumeDatabaseDriver, " does not support ", name, ".")
			errors = true
		}
	}

	noteErrorFunc("-dbpath", volumeDatabasePath, path)
	noteErrorFunc("-dbhost", volumeDatabaseHost, host)
	noteErrorFunc("-dbuser", volumeDatabaseUsername, username)
	noteErrorFunc("-dbpass", volumeDatabasePassword, password)
	noteErrorFunc("-dbschema", volumeDatabaseSchema, schema)

	if errors && fatal {
		glog.Fatal("Invalid flag(s) were passed, are you using the correct volume driver?")
	}
}

func getStorageConnection() (drivers.StorageController, error) {
	glog.Info("Attempting to use the ", storageControllerDriver, " storage controller.")

	switch storageControllerDriver {
	case "on-disk":
		validateStorageFlags(true, true)
		return drivers.NewOnDiskStorageController(storageControllerPath), nil

	case "glusterfs":
		validateStorageFlags(false, false)
		return drivers.NewGlusterStorageController(), nil

	default:
		return nil, errors.New("unsupported storage controller, please choose glusterfs or on-disk")
	}
}

func validateStorageFlags(fatal bool, path bool) {

	var errors bool
	noteErrorFunc := func(name string, value string, used bool) {
		if !used && value != "" {
			glog.Warning("Storage Controller: ", volumeDatabaseDriver, " does not support ", name, ".")
			errors = true
		}
	}

	noteErrorFunc("-dbpath", storageControllerPath, path)

	if errors && fatal {
		glog.Fatal("Invalid flag(s) were passed, are you using the correct storage controller?")
	}
}
<commit_msg>Add configurable plugin name to launch multiple instances on the same machine<commit_after>// Creates a webserver/unix socket that allows a Docker server to create RDMA
// backed volumes.
package main

import (
	"errors"
	"flag"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/docker/go-plugins-helpers/volume"
	"github.com/golang/glog"
	"github.com/mellanox-senior-design/docker-volume-rdma/db"
	"github.com/mellanox-senior-design/docker-volume-rdma/drivers"
)

// Name of the plugin for use in Docker CLI
var pluginName string

// Port to launch service on.
var httpPort int

// Volume Flags
var volumeDatabaseDriver string
var volumeDatabasePath string
var volumeDatabaseHost string
var volumeDatabaseUsername string
var volumeDatabasePassword string
var volumeDatabaseSchema string

// Storage Flags
var storageControllerDriver string
var storageControllerPath string

func init() {
	// Configure application flags.
	flag.StringVar(&pluginName, "name", "docker-volume-rdma", "name of the plugin used in the Docker CLI")
	flag.IntVar(&httpPort, "port", 8080, "tcp/ip port to serve volume driver on")

	// Volume Database Flags
	flag.StringVar(&volumeDatabaseDriver, "db", "sqlite", "set the database backend used to store volume metadata: [sqlite, mysql, in-memory]")
	flag.StringVar(&volumeDatabasePath, "dbpath", "", "set the database storage path")
	flag.StringVar(&volumeDatabaseHost, "dbhost", "", "set the database host (default is '' or localhost:3306)")
	flag.StringVar(&volumeDatabaseUsername, "dbuser", "", "set the database username (default is root)")
	flag.StringVar(&volumeDatabasePassword, "dbpass", "", "set the database password (optional)")
	flag.StringVar(&volumeDatabaseSchema, "dbschema", "", "set the database schema (required)")

	// Storage Controller Flags
	flag.StringVar(&storageControllerDriver, "sc", "glusterfs", "set the storage backend used to store volume data: [glusterfs, on-disk]")
	flag.StringVar(&storageControllerPath, "scpath", "", "set the storage path used to know where to put the volumes on the host")
}

// Configure and start the docker volume plugin server.
func main() {

	// Parse flags as glog needs the flags to be solidified before starting.
	flag.Parse()

	// Convert port to string, and print startup message.
	port := strconv.Itoa(httpPort)

	driver, handler, err := configure()
	if err == nil {

		// Handle SIGINT gracefully
		sigintChannel := make(chan os.Signal, 2)
		signal.Notify(sigintChannel, os.Interrupt, syscall.SIGINT)
		go func() {
			<-sigintChannel
			glog.Info("Exiting ...")
			driver.Disconnect()
			os.Exit(1)
		}()

		err = driver.Connect()
		if err == nil {
			defer driver.Disconnect()

			glog.Info("Running! http://localhost:" + port)
			err = handler.ServeTCP(pluginName, ":"+port, nil)
		}
	}

	// Log any error that may have occurred.
	glog.Fatal(err)
}

func configure() (*drivers.RDMAVolumeDriver, *volume.Handler, error) {
	// Create and begin serving volume driver on tcp/ip port, httpPort.
	volumeDatabase, err := getDatabaseConnection()
	if err != nil {
		return nil, nil, err
	}

	// Configure Storage Controller
	storageController, err := getStorageConnection()
	if err != nil {
		return nil, nil, err
	}

	// Print startup message and start server
	glog.Info("Connecting to services ...")
	driver := drivers.NewRDMAVolumeDriver(storageController, volumeDatabase)
	handler := volume.NewHandler(driver)

	return &driver, handler, nil
}

// GetDatabaseConnection returns the database connection that was requested on the command line.
func getDatabaseConnection() (db.VolumeDatabase, error) {
	glog.Info("Attempting to use the ", volumeDatabaseDriver, " volume driver.")
	switch volumeDatabaseDriver {
	case "in-memory":
		validateDatabaseFlags(true, true, false, false, false, false)
		return db.NewInMemoryVolumeDatabase(), nil

	case "sqlite":
		validateDatabaseFlags(false, true, false, false, false, false)
		return db.NewSQLiteVolumeDatabase(volumeDatabasePath), nil

	case "mysql":
		validateDatabaseFlags(false, false, true, true, true, true)
		return db.NewMySQLVolumeDatabase(volumeDatabaseHost, volumeDatabaseUsername, volumeDatabasePassword, volumeDatabaseSchema)

	default:
		return nil, errors.New("unsupported database, please choose sqlite, mysql, or in-memory")
	}
}

func validateDatabaseFlags(fatal bool, path bool, host bool, username bool, password bool, schema bool) {

	var errors bool
	noteErrorFunc := func(name string, value string, used bool) {
		if !used && value != "" {
			glog.Warning("Volume Driver: ", volumeDatabaseDriver, " does not support ", name, ".")
			errors = true
		}
	}

	noteErrorFunc("-dbpath", volumeDatabasePath, path)
	noteErrorFunc("-dbhost", volumeDatabaseHost, host)
	noteErrorFunc("-dbuser", volumeDatabaseUsername, username)
	noteErrorFunc("-dbpass", volumeDatabasePassword, password)
	noteErrorFunc("-dbschema", volumeDatabaseSchema, schema)

	if errors && fatal {
		glog.Fatal("Invalid flag(s) were passed, are you using the correct volume driver?")
	}
}

func getStorageConnection() (drivers.StorageController, error) {
	glog.Info("Attempting to use the ", storageControllerDriver, " storage controller.")

	switch storageControllerDriver {
	case "on-disk":
		validateStorageFlags(true, true)
		return drivers.NewOnDiskStorageController(storageControllerPath), nil

	case "glusterfs":
		validateStorageFlags(false, false)
		return drivers.NewGlusterStorageController(), nil

	default:
		return nil, errors.New("unsupported storage controller, please choose glusterfs or on-disk")
	}
}

func validateStorageFlags(fatal bool, path bool) {

	var errors bool
	noteErrorFunc := func(name string, value string, used bool) {
		if !used && value != "" {
			glog.Warning("Storage Controller: ", volumeDatabaseDriver, " does not support ", name, ".")
			errors = true
		}
	}

	noteErrorFunc("-dbpath", storageControllerPath, path)

	if errors && fatal {
		glog.Fatal("Invalid flag(s) were passed, are you using the correct storage controller?")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"runtime"
	"time"

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

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

type scanConfig struct {
	cores int
}

type appConfig struct {
	printUrls bool
}

type config struct {
	app  appConfig
	scan scanConfig
	out  outputConfig
}

var conf config
var out = Output{}

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

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

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

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

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

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

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

	return
}

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

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

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

	return nil
}

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

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

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

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

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

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

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

	app := cli.NewApp()

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

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

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

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

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

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

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

			os.Exit(0)
		}

		run()

		return nil
	}

	app.Run(os.Args)
}
<commit_msg>use placeholders<commit_after>package main

import (
	"fmt"
	"os"
	"runtime"
	"time"

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

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

type scanConfig struct {
	cores int
}

type appConfig struct {
	printUrls bool
}

type config struct {
	app  appConfig
	scan scanConfig
	out  outputConfig
}

var conf config
var out = Output{}

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

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

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

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

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

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

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

	return
}

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

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

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

	return nil
}

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

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

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

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

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

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

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

	app := cli.NewApp()

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

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

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

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

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

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

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

			os.Exit(0)
		}

		run()

		return nil
	}

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


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"
)

type Acrophone struct {
	charMap map[string]string
}

func (a *Acrophone) Lookup(char string) (string, error) {
	if len(char) > 1 {
		return "", fmt.Errorf("Lookup() accepts single character argument only")
	}
	if res, ok := a.charMap[char]; ok == false {
		return "", fmt.Errorf("unable to find match for character: %s", char)
	} else {
		return res, nil
	}
}

func main() {
	var output []string

	if len(os.Args) < 2 {
		fmt.Println("usage: acrophone <input>")
		os.Exit(1)
	}

	a := &Acrophone{NatoMap}

	input := strings.Join(os.Args[1:], "")
	for _, char := range input {
		result, err := a.Lookup(string(char))
		if err != nil {
			panic(err)
		}
		output = append(output, result)
	}

	fmt.Println(strings.Join(output, " "))
}
<commit_msg>non-case-sensitive lookup<commit_after>package main

import (
	"fmt"
	"os"
	"strings"
)

type Acrophone struct {
	charMap map[string]string
}

func (a *Acrophone) Lookup(char string) (string, error) {
	char = strings.ToLower(char)

	if len(char) > 1 {
		return "", fmt.Errorf("Lookup() accepts single character argument only")
	}

	if res, ok := a.charMap[char]; ok == false {
		return "", fmt.Errorf("unable to find match for character: %s", char)
	} else {
		return res, nil
	}
}

func main() {
	var output []string

	if len(os.Args) < 2 {
		fmt.Println("usage: acrophone <input>")
		os.Exit(1)
	}

	a := &Acrophone{NatoMap}

	input := strings.Join(os.Args[1:], "")
	for _, char := range input {
		result, err := a.Lookup(string(char))
		if err != nil {
			panic(err)
		}
		output = append(output, result)
	}

	fmt.Println(strings.Join(output, " "))
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"

	_ "github.com/mattn/go-sqlite3"
)

var wd string

func init() {
	workingDir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	wd = workingDir
}

func main() {
	b := &blog{}
	b.load("blog.cfg")
}
<commit_msg>blog: split load and start functions<commit_after>package main

import (
	"log"
	"os"

	_ "github.com/mattn/go-sqlite3"
)

var wd string

func init() {
	workingDir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	wd = workingDir
}

func main() {
	b := &blog{}
	b.load("blog.cfg")
	b.start()

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/luismesas/goPi/piface"
	"github.com/luismesas/goPi/spi"
	"time"
)

func main() {

	// creates a new pifacedigital instance
	pfd := piface.NewPiFaceDigital(spi.DEFAULT_HARDWARE_ADDR, spi.DEFAULT_BUS, spi.DEFAULT_CHIP)

	// initializes pifacedigital board
	err := pfd.InitBoard()
	if err != nil {
		fmt.Printf("Error on init board: %s", err)
		return
	}

	// blink time!!
	fmt.Println("Blinking led 7 each second")
	pfd.Leds[6].Toggle()
	for {
		pfd.Leds[7].Toggle()
		pfd.Leds[6].Toggle()
		time.Sleep(time.Second)
	}
}
<commit_msg>all on off<commit_after>package main

import (
	"fmt"
	"github.com/luismesas/goPi/piface"
	"github.com/luismesas/goPi/spi"
	"time"
)

func main() {

	// creates a new pifacedigital instance
	pfd := piface.NewPiFaceDigital(spi.DEFAULT_HARDWARE_ADDR, spi.DEFAULT_BUS, spi.DEFAULT_CHIP)

	// initializes pifacedigital board
	err := pfd.InitBoard()
	if err != nil {
		fmt.Printf("Error on init board: %s", err)
		return
	}

	for k, v := range pfd.Leds {
		v.AllOn()
		time.Sleep(time.Second)
	}
	for k, v = range pfd.Leds {
		v.AllOff()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"sync"
)

const (
	// Timeout for DNS queries
	timeout = 3 * 1e9

	// maximum number of attempts for a query
	maxAttempts = 3
)

var (
	pending         = make(chan *job, 100)
	finished        = make(chan *job, 100)
	done            sync.WaitGroup
	workersCount    = 32
	referenceServer = "8.8.8.8"
	connection      string
	domainArg       string
)

func main() {
	databaseArg := flag.String("database", "database.yml", "Path to file containing the database configuration")
	flag.StringVar(&domainArg, "domains", "domains.txt", "Path to file containing the domain list")
	flag.StringVar(&geoDbPath, "geodb", "GeoLite2-City.mmdb", "Path to GeoDB database")
	flag.StringVar(&referenceServer, "reference", referenceServer, "The nameserver that every other is compared with")
	flag.IntVar(&workersCount, "workers", workersCount, "Number of worker routines")
	flag.Parse()

	dnsClient.ReadTimeout = timeout

	environment := os.Getenv("RAILS_ENV")
	if environment == "" {
		environment = "development"
	}

	// read domain list
	if err := readDomains(domainArg); err != nil {
		fmt.Println("unable to read domain list")
		panic(err)
	}

	// load database configuration
	connection = databasePath(*databaseArg, environment)

	// check the GeoDB
	location(referenceServer)

	// Use all cores

	// Get results from the reference nameserver
	res, _, err := resolveDomains(referenceServer)
	if err != nil {
		panic(err)
	}
	expectedResults = res

	// Start result writer
	go resultWriter()

	// Start workers
	done.Add(workersCount)
	for i := 0; i < workersCount; i++ {
		go worker()
	}

	createJobs()

	// wait for workers to finish
	done.Wait()

	close(finished)
}

func createJobs() {
	currentId := 0
	batchSize := 1000
	found := batchSize

	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	for batchSize == found {
		// Read the next batch
		rows, err := db.Query("SELECT id, ip FROM nameservers WHERE id > ? LIMIT ?", currentId, batchSize)
		if err != nil {
			panic(err)
		}

		found = 0
		for rows.Next() {
			j := new(job)

			// get RawBytes from data
			err = rows.Scan(&j.id, &j.address)
			if err != nil {
				panic(err)
			}
			pending <- j
			currentId = j.id
			found += 1
		}
		rows.Close()
	}
	close(pending)
}

func worker() {
	for job := range pending {
		executeJob(job)
		finished <- job
	}
	done.Done()
}

func resultWriter() {
	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	stm, err := db.Prepare(
		"UPDATE nameservers SET name=?, state=?, error=?, version=?, dnssec=?, checked_at=NOW(), country_id=?, city=?," +
			"state_changed_at = (CASE WHEN ? != state THEN NOW() ELSE state_changed_at END )" +
			"WHERE id=?")
	defer stm.Close()

	for res := range finished {
		log.Println(res)
		stm.Exec(res.name, res.state, res.err, res.version, res.dnssec, res.country, res.city, res.state, res.id)
	}
}

// consumes a job and writes the result in the given job
func executeJob(job *job) {
	// log.Println("received job", job.id)

	// GeoDB lookup
	job.country, job.city = location(job.address)

	// Run the check
	err, dnssec := check(job)
	job.name = ptrName(job.address)

	// query the bind version
	if err == nil || err.Error() != "i/o timeout" {
		job.version = version(job.address)
	}

	if err == nil {
		job.state = "valid"
		job.err = ""
		job.dnssec = &dnssec
	} else {
		job.state = "invalid"
		job.err = err.Error()
	}
}
<commit_msg>Wait for the resultWriter to finish<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"sync"
)

const (
	// Timeout for DNS queries
	timeout = 3 * 1e9

	// maximum number of attempts for a query
	maxAttempts = 3
)

var (
	pending         = make(chan *job, 100)
	finished        = make(chan *job, 100)
	pendingWg       sync.WaitGroup
	finishedWg      sync.WaitGroup
	workersCount    = 32
	referenceServer = "8.8.8.8"
	connection      string
	domainArg       string
)

func main() {
	databaseArg := flag.String("database", "database.yml", "Path to file containing the database configuration")
	flag.StringVar(&domainArg, "domains", "domains.txt", "Path to file containing the domain list")
	flag.StringVar(&geoDbPath, "geodb", "GeoLite2-City.mmdb", "Path to GeoDB database")
	flag.StringVar(&referenceServer, "reference", referenceServer, "The nameserver that every other is compared with")
	flag.IntVar(&workersCount, "workers", workersCount, "Number of worker routines")
	flag.Parse()

	dnsClient.ReadTimeout = timeout

	environment := os.Getenv("RAILS_ENV")
	if environment == "" {
		environment = "development"
	}

	// read domain list
	if err := readDomains(domainArg); err != nil {
		fmt.Println("unable to read domain list")
		panic(err)
	}

	// load database configuration
	connection = databasePath(*databaseArg, environment)

	// check the GeoDB
	location(referenceServer)

	// Use all cores

	// Get results from the reference nameserver
	res, _, err := resolveDomains(referenceServer)
	if err != nil {
		panic(err)
	}
	expectedResults = res

	// Start result writer
	finishedWg.Add(1)
	go resultWriter()

	// Start workers
	pendingWg.Add(workersCount)
	for i := 0; i < workersCount; i++ {
		go worker()
	}

	createJobs()

	// wait for workers to finish
	pendingWg.Wait()

	close(finished)
	finishedWg.Wait()
}

func createJobs() {
	currentId := 0
	batchSize := 1000
	found := batchSize

	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	for batchSize == found {
		// Read the next batch
		rows, err := db.Query("SELECT id, ip FROM nameservers WHERE id > ? LIMIT ?", currentId, batchSize)
		if err != nil {
			panic(err)
		}

		found = 0
		for rows.Next() {
			j := new(job)

			// get RawBytes from data
			err = rows.Scan(&j.id, &j.address)
			if err != nil {
				panic(err)
			}
			pending <- j
			currentId = j.id
			found += 1
		}
		rows.Close()
	}
	close(pending)
}

func worker() {
	for job := range pending {
		executeJob(job)
		finished <- job
	}
	pendingWg.Done()
}

func resultWriter() {
	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	stm, err := db.Prepare(
		"UPDATE nameservers SET name=?, state=?, error=?, version=?, dnssec=?, checked_at=NOW(), country_id=?, city=?," +
			"state_changed_at = (CASE WHEN ? != state THEN NOW() ELSE state_changed_at END )" +
			"WHERE id=?")
	defer stm.Close()

	for res := range finished {
		log.Println(res)
		stm.Exec(res.name, res.state, res.err, res.version, res.dnssec, res.country, res.city, res.state, res.id)
	}

	finishedWg.Done()
}

// consumes a job and writes the result in the given job
func executeJob(job *job) {
	// log.Println("received job", job.id)

	// GeoDB lookup
	job.country, job.city = location(job.address)

	// Run the check
	err, dnssec := check(job)
	job.name = ptrName(job.address)

	// query the bind version
	if err == nil || err.Error() != "i/o timeout" {
		job.version = version(job.address)
	}

	if err == nil {
		job.state = "valid"
		job.err = ""
		job.dnssec = &dnssec
	} else {
		job.state = "invalid"
		job.err = err.Error()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/ogier/pflag"
	"github.com/yuya-takeyama/argf"
)

func printUsage() {
	os.Stderr.WriteString(`
Usage: csvp [OPTION]... [FILE]...
Print selected parts of CSV from each FILE to standard output.

Options:
  -i, --indexes=LIST
                 select only these indexes
  -h, --headers=LIST
                 select only these headers
  -t, --tsv
                 equivalent to -d'\t'
  -d, --delimiter=DELIM
                 use DELIM instead of comma for field delimiter
  -D, --output-delimiter=STRING
                 use STRING as the output delimiter (default: \t)
  --help
                 display this help text and exit
  --version
                 output version information and exit
`[1:])
}

func printVersion() {
	os.Stderr.WriteString(`
0.8.1
`[1:])
}

type Option struct {
	IndexesList     string
	HeadersList     string
	IsTSV           bool
	Delimiter       string
	OutputDelimiter string
	IsHelp          bool
	IsVersion       bool
	Files           []string
}

func parseOption(args []string) (opt *Option, err error) {
	flag := pflag.NewFlagSet("csvp", pflag.ContinueOnError)
	flag.SetOutput(ioutil.Discard)

	opt = &Option{}
	flag.StringVarP(&opt.IndexesList, "indexes", "i", "", "")
	flag.StringVarP(&opt.HeadersList, "headers", "h", "", "")
	flag.BoolVarP(&opt.IsTSV, "tsv", "t", false, "")
	flag.StringVarP(&opt.Delimiter, "delimiter", "d", ",", "")
	flag.StringVarP(&opt.OutputDelimiter, "output-delimiter", "D", "\t", "")
	flag.BoolVarP(&opt.IsHelp, "help", "", false, "")
	flag.BoolVarP(&opt.IsVersion, "version", "", false, "")

	if err = flag.Parse(args); err != nil {
		return nil, err
	}

	opt.Files = flag.Args()
	return opt, nil
}

func toDelimiter(s string) (r rune, err error) {
	s, err = strconv.Unquote(`"` + s + `"`)
	if err != nil {
		return 0, err
	}

	runes := []rune(s)
	if len(runes) != 1 {
		return 0, fmt.Errorf("the delimiter must be a single character")
	}
	return runes[0], nil
}

func newCSVScannerFromOption(opt *Option) (c *CSVScanner, err error) {
	var selector Selector
	switch {
	case opt.IndexesList != "" && opt.HeadersList != "":
		return nil, fmt.Errorf("only one type of list may be specified")
	case opt.IndexesList != "":
		selector = NewIndexes(opt.IndexesList)
	case opt.HeadersList != "":
		selector = NewHeaders(opt.HeadersList)
	default:
		selector = NewAll()
	}

	reader, err := argf.From(opt.Files)
	if err != nil {
		return nil, err
	}

	c = NewCSVScanner(selector, reader)
	c.SetOutputDelimiter(opt.OutputDelimiter)
	switch {
	case opt.IsTSV:
		c.SetDelimiter('\t')
	default:
		delimiter, err := toDelimiter(opt.Delimiter)
		if err != nil {
			return nil, err
		}
		c.SetDelimiter(delimiter)
	}
	return c, nil
}

func do(c *CSVScanner) error {
	for c.Scan() {
		fmt.Println(c.Text())
	}
	return c.Err()
}

func printErr(err error) {
	fmt.Fprintln(os.Stderr, "cspv:", err)
}

func guideToHelp() {
	os.Stderr.WriteString(`
Try 'cspv --help' for more information.
`[1:])
}

func _main() int {
	opt, err := parseOption(os.Args[1:])
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	switch {
	case opt.IsHelp:
		printUsage()
		return 0
	case opt.IsVersion:
		printVersion()
		return 0
	}

	c, err := newCSVScannerFromOption(opt)
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	if err = do(c); err != nil {
		printErr(err)
		return 1
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<commit_msg>Extract version<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/ogier/pflag"
	"github.com/yuya-takeyama/argf"
)

var (
	version = "0.8.1"
)

func printUsage() {
	os.Stderr.WriteString(`
Usage: csvp [OPTION]... [FILE]...
Print selected parts of CSV from each FILE to standard output.

Options:
  -i, --indexes=LIST
                 select only these indexes
  -h, --headers=LIST
                 select only these headers
  -t, --tsv
                 equivalent to -d'\t'
  -d, --delimiter=DELIM
                 use DELIM instead of comma for field delimiter
  -D, --output-delimiter=STRING
                 use STRING as the output delimiter (default: \t)
  --help
                 display this help text and exit
  --version
                 output version information and exit
`[1:])
}

func printVersion() {
	fmt.Fprintln(os.Stderr, version)
}

type Option struct {
	IndexesList     string
	HeadersList     string
	IsTSV           bool
	Delimiter       string
	OutputDelimiter string
	IsHelp          bool
	IsVersion       bool
	Files           []string
}

func parseOption(args []string) (opt *Option, err error) {
	flag := pflag.NewFlagSet("csvp", pflag.ContinueOnError)
	flag.SetOutput(ioutil.Discard)

	opt = &Option{}
	flag.StringVarP(&opt.IndexesList, "indexes", "i", "", "")
	flag.StringVarP(&opt.HeadersList, "headers", "h", "", "")
	flag.BoolVarP(&opt.IsTSV, "tsv", "t", false, "")
	flag.StringVarP(&opt.Delimiter, "delimiter", "d", ",", "")
	flag.StringVarP(&opt.OutputDelimiter, "output-delimiter", "D", "\t", "")
	flag.BoolVarP(&opt.IsHelp, "help", "", false, "")
	flag.BoolVarP(&opt.IsVersion, "version", "", false, "")

	if err = flag.Parse(args); err != nil {
		return nil, err
	}

	opt.Files = flag.Args()
	return opt, nil
}

func toDelimiter(s string) (r rune, err error) {
	s, err = strconv.Unquote(`"` + s + `"`)
	if err != nil {
		return 0, err
	}

	runes := []rune(s)
	if len(runes) != 1 {
		return 0, fmt.Errorf("the delimiter must be a single character")
	}
	return runes[0], nil
}

func newCSVScannerFromOption(opt *Option) (c *CSVScanner, err error) {
	var selector Selector
	switch {
	case opt.IndexesList != "" && opt.HeadersList != "":
		return nil, fmt.Errorf("only one type of list may be specified")
	case opt.IndexesList != "":
		selector = NewIndexes(opt.IndexesList)
	case opt.HeadersList != "":
		selector = NewHeaders(opt.HeadersList)
	default:
		selector = NewAll()
	}

	reader, err := argf.From(opt.Files)
	if err != nil {
		return nil, err
	}

	c = NewCSVScanner(selector, reader)
	c.SetOutputDelimiter(opt.OutputDelimiter)
	switch {
	case opt.IsTSV:
		c.SetDelimiter('\t')
	default:
		delimiter, err := toDelimiter(opt.Delimiter)
		if err != nil {
			return nil, err
		}
		c.SetDelimiter(delimiter)
	}
	return c, nil
}

func do(c *CSVScanner) error {
	for c.Scan() {
		fmt.Println(c.Text())
	}
	return c.Err()
}

func printErr(err error) {
	fmt.Fprintln(os.Stderr, "cspv:", err)
}

func guideToHelp() {
	os.Stderr.WriteString(`
Try 'cspv --help' for more information.
`[1:])
}

func _main() int {
	opt, err := parseOption(os.Args[1:])
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	switch {
	case opt.IsHelp:
		printUsage()
		return 0
	case opt.IsVersion:
		printVersion()
		return 0
	}

	c, err := newCSVScannerFromOption(opt)
	if err != nil {
		printErr(err)
		guideToHelp()
		return 2
	}
	if err = do(c); err != nil {
		printErr(err)
		return 1
	}
	return 0
}

func main() {
	e := _main()
	os.Exit(e)
}
<|endoftext|>