<commit_before>package main

import (
	"time"
	"encoding/json"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"log"
	"os"
	"io"
	"fmt"
	"database/sql"
	_ "github.com/lib/pq"
)

var conn *sql.DB

func SetHeaders(w *http.ResponseWriter) {
	(*w).Header().Set("Access-Control-Allow-Origin", "*")
	(*w).Header().Set("Content-Type", "application/json")
}

func GetBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM blog")
	data := []BlogPost{}
	for rows.Next() {
		post := BlogPost{}
		rows.Scan(&post.Id, &post.Titel, &post.Text, &post.Auteur, &post.Img_url, &post.Ctime, &post.Image)
		data = append(data, post)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata ORDER BY tijd DESC LIMIT 1")
	spin := SpinData{}
	rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
	buf,_ := json.Marshal(spin)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata")
	data := []SpinData{}
	for rows.Next() {
		spin := SpinData{}
		rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
		data = append(data, spin)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT batterij FROM spindata")
	data := make([]int, 0)
	var scanInt int
	for rows.Next() {
		rows.Scan(&scanInt)
		data = append(data, scanInt)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata ORDER BY tijd DESC LIMIT 1")
	servo := ServoData{}
	rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
	buf,_ := json.Marshal(servo)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata")
	data := []ServoData{}
	for rows.Next() {
		servo := ServoData{}
		rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
		data = append(data, servo)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func PostBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	r.ParseMultipartForm(32 << 20)
	file, handler, err := r.FormFile("uploadfile")
	if err != nil {
		w.WriteHeader(500)
		return
	}
	defer file.Close()
	fmt.Fprintf(w, "%v", handler.Header)
	f, err := os.OpenFile("./img/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0666)
	if err != nil {
		w.WriteHeader(500)
		return
	}
	defer f.Close()
	io.Copy(f, file)
	_,err = conn.Query("INSERT INTO blog (titel, text, auteur, ctime, image) VALUES ($1, $2, $3, $4, $5)", r.FormValue("titel"), r.FormValue("text"), r.FormValue("auteur"), time.Now(), "http://idp-api.herokuapp.com/img/"+handler.Filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("error posting"))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO spindata (tijd, mode, hellingsgraad, snelheid, batterij, balloncount) VALUES ($1, $2, $3, $4, $5, $6)", time.Now(),
		r.FormValue("mode"), r.FormValue("hellingsgraad"), r.FormValue("snelheid"), r.FormValue("batterij"), r.FormValue("ballonCount"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO servodata (servo_id, tijd, voltage, positie, load, temperatuur) VALUES ($1, $2, $3, $4, $5, $6)",
		r.FormValue("servo_id"), time.Now(), r.FormValue("voltage"), r.FormValue("positie"), r.FormValue("load"), r.FormValue("Temperatuur"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func GetImg(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	http.ServeFile(w,r,"./img/"+ps.ByName("file"))
}

func GetDoc(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	http.ServeFile(w,r,"./doc.html")
}

func main() {
	conn,_ = sql.Open("postgres", os.Getenv("DATABASE_URL"))
	defer conn.Close()
	router := httprouter.New()
	router.GET("/", GetDoc)
	router.GET("/blog", GetBlog)
	router.GET("/spin/latest", GetLatestSpinData)
	router.GET("/spin/archive", GetArchivedSpinData)
	router.GET("/spin/archive/batterij", GetArchivedSpinBatterij)
	router.GET("/servo/latest", GetLatestServoData)
	router.GET("/servo/archive", GetArchivedServoData)
	router.GET("/img/*file", GetImg)
	router.POST("/blog", PostBlog)
	router.POST("/spin", PostSpinData)
	router.POST("/servo", PostServoData)
	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}
	fmt.Printf("Starting server at localhost:%s...", port)
	log.Fatal(http.ListenAndServe(":"+port, router))
}<commit_msg>Archived mode route en functie<commit_after>package main

import (
	"time"
	"encoding/json"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"log"
	"os"
	"io"
	"fmt"
	"database/sql"
	_ "github.com/lib/pq"
)

var conn *sql.DB

func SetHeaders(w *http.ResponseWriter) {
	(*w).Header().Set("Access-Control-Allow-Origin", "*")
	(*w).Header().Set("Content-Type", "application/json")
}

func GetBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM blog")
	data := []BlogPost{}
	for rows.Next() {
		post := BlogPost{}
		rows.Scan(&post.Id, &post.Titel, &post.Text, &post.Auteur, &post.Img_url, &post.Ctime, &post.Image)
		data = append(data, post)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata ORDER BY tijd DESC LIMIT 1")
	spin := SpinData{}
	rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
	buf,_ := json.Marshal(spin)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM spindata")
	data := []SpinData{}
	for rows.Next() {
		spin := SpinData{}
		rows.Scan(&spin.Id, &spin.Tijd, &spin.Mode, &spin.Hellingsgraad, &spin.Snelheid, &spin.Batterij, &spin.BallonCount)
		data = append(data, spin)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinBatterij(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT batterij FROM spindata")
	data := make([]int, 0)
	var scanInt int
	for rows.Next() {
		rows.Scan(&scanInt)
		data = append(data, scanInt)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedSpinMode(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT mode FROM spindata")
	data := make([]string, 0)
	var scanStr string
	for rows.Next() {
		rows.Scan(&scanStr)
		data = append(data, scanStr)
	}
	buf,_ := json.Marshal(data)
	fmt.Printf(string(buf))
	SetHeaders(&w)
	w.Write(buf)
}

func GetLatestServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata ORDER BY tijd DESC LIMIT 1")
	servo := ServoData{}
	rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
	buf,_ := json.Marshal(servo)
	SetHeaders(&w)
	w.Write(buf)
}

func GetArchivedServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rows,_ := conn.Query("SELECT * FROM servodata")
	data := []ServoData{}
	for rows.Next() {
		servo := ServoData{}
		rows.Scan(&servo.Id, &servo.ServoId, &servo.Tijd, &servo.Voltage, &servo.Positie, &servo.Load, &servo.Temperatuur)
		data = append(data, servo)
	}
	buf,_ := json.Marshal(data)
	SetHeaders(&w)
	w.Write(buf)
}

func PostBlog(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	r.ParseMultipartForm(32 << 20)
	file, handler, err := r.FormFile("uploadfile")
	if err != nil {
		w.WriteHeader(500)
		return
	}
	defer file.Close()
	fmt.Fprintf(w, "%v", handler.Header)
	f, err := os.OpenFile("./img/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0666)
	if err != nil {
		w.WriteHeader(500)
		return
	}
	defer f.Close()
	io.Copy(f, file)
	_,err = conn.Query("INSERT INTO blog (titel, text, auteur, ctime, image) VALUES ($1, $2, $3, $4, $5)", r.FormValue("titel"), r.FormValue("text"), r.FormValue("auteur"), time.Now(), "http://idp-api.herokuapp.com/img/"+handler.Filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("error posting"))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostSpinData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO spindata (tijd, mode, hellingsgraad, snelheid, batterij, balloncount) VALUES ($1, $2, $3, $4, $5, $6)", time.Now(),
		r.FormValue("mode"), r.FormValue("hellingsgraad"), r.FormValue("snelheid"), r.FormValue("batterij"), r.FormValue("ballonCount"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func PostServoData(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	_,err := conn.Query("INSERT INTO servodata (servo_id, tijd, voltage, positie, load, temperatuur) VALUES ($1, $2, $3, $4, $5, $6)",
		r.FormValue("servo_id"), time.Now(), r.FormValue("voltage"), r.FormValue("positie"), r.FormValue("load"), r.FormValue("Temperatuur"))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.WriteHeader(201)
	w.Write([]byte("successful"))
}

func GetImg(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	http.ServeFile(w,r,"./img/"+ps.ByName("file"))
}

func GetDoc(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	http.ServeFile(w,r,"./doc.html")
}

func main() {
	conn,_ = sql.Open("postgres", os.Getenv("DATABASE_URL"))
	defer conn.Close()
	router := httprouter.New()
	router.GET("/", GetDoc)
	router.GET("/blog", GetBlog)
	router.GET("/spin/latest", GetLatestSpinData)
	router.GET("/spin/archive", GetArchivedSpinData)
	router.GET("/spin/archive/batterij", GetArchivedSpinBatterij)
	router.GET("/spin/archive/mode", GetArchivedSpinMode)
	router.GET("/servo/latest", GetLatestServoData)
	router.GET("/servo/archive", GetArchivedServoData)
	router.GET("/img/*file", GetImg)
	router.POST("/blog", PostBlog)
	router.POST("/spin", PostSpinData)
	router.POST("/servo", PostServoData)
	port := os.Getenv("PORT")
	if port == "" {
		port = "3000"
	}
	fmt.Printf("Starting server at localhost:%s...", port)
	log.Fatal(http.ListenAndServe(":"+port, router))
}<|endoftext|>


<commit_before>package main

import (
	"log"

	"github.com/fsufitch/prez-tweet/prez-tweet-server"
)

func main() {
	log.Fatal(server.StartServer())
}
<commit_msg>Add interrupt signal processing<commit_after>package main

import (
	"log"
	"os"
	"os/signal"

	"github.com/fsufitch/prez-tweet/prez-tweet-server"
)

func main() {
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	serverQuitChan := make(chan error)
	go func() {
		serverQuitChan <- server.StartServer()
	}()

	select {
	case <-interruptChan:
		log.Println("Received interrupt. Quitting.")
		break
	case err := <-serverQuitChan:
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// 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 (
	"fmt"
	"log"
	"net/http"
	"os"
	"net/url"
	"io/ioutil"

	"github.com/line/line-bot-sdk-go/linebot"
)

var bot *linebot.Client

func main() {
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	events, err := bot.ParseRequest(r)

	resp, err := http.PostForm("https://hr.hiyes.tw:443/getMessage.php",
        url.Values{"mid": {"ziv"}, "message": {"test"}})
    if err != nil {
        fmt.Println(err)
    } else {
        body, _ := ioutil.ReadAll(resp.Body)
        fmt.Println("POST OK: ", string(body), resp)
    }

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {
			switch message := event.Message.(type) {
			case *linebot.TextMessage:
				if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(event.ReplyToken+"---"+message.Text+" OK!")).Do(); err != nil {
					log.Print(err)
				}
			}
		}
	}
}
<commit_msg>edit reply msg<commit_after>// 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 (
	"fmt"
	"log"
	"net/http"
	"os"
	"net/url"
	"io/ioutil"

	"github.com/line/line-bot-sdk-go/linebot"
)

var bot *linebot.Client

func main() {
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	events, err := bot.ParseRequest(r)

	resp, err := http.PostForm("https://hr.hiyes.tw:443/getMessage.php",
        url.Values{"mid": {"ziv"}, "message": {"test"}})
    if err != nil {
        fmt.Println(err)
    } else {
        body, _ := ioutil.ReadAll(resp.Body)
        fmt.Println("POST OK: ", string(body), resp)
    }

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {
			switch message := event.Message.(type) {
			case *linebot.TextMessage:
				if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(message.id+"---"+message.Text+" OK!")).Do(); err != nil {
					log.Print(err)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(3)

	producerSequence := NewSequence()
	consumerSequence1 := NewSequence()
	consumerSequence2 := NewSequence()

	producerBarrier := NewBarrier(producerSequence)
	consumerBarrier := NewBarrier(consumerSequence1, consumerSequence2)

	sequencer := NewSingleProducerSequencer(producerSequence, RingSize, consumerBarrier)
	go consume(producerBarrier, producerSequence, consumerSequence1)
	go consume(producerBarrier, producerSequence, consumerSequence2)

	started := time.Now()
	for i := int64(0); i < MaxSequenceValue; i++ {
		ticket := sequencer.Next(1)
		ringBuffer[ticket&RingMask] = ticket
		sequencer.Publish(ticket)
		if ticket%Mod == 0 && ticket > 0 {
			finished := time.Now()
			elapsed := finished.Sub(started)
			fmt.Println(ticket, elapsed)
			started = time.Now()
		}
	}
}

func consume(barrier Barrier, source, sequence *Sequence) {
	worker := NewWorker(barrier, TestHandler{}, source, sequence)

	for {
		worker.Process()
	}
}

const Mod = 1000000 * 100 // 1 million * 10
const RingSize = 1024 * 256
const RingMask = RingSize - 1

var ringBuffer [RingSize]int64

type TestHandler struct{}

func (this TestHandler) Consume(sequence, remaining int64) {
	message := ringBuffer[sequence&RingMask]
	if message != sequence {
		fmt.Printf("ERROR Consumer:: Sequence: %d, Message: %d\n", sequence, message)
		panic(fmt.Sprintf("Consumer:: Sequence: %d, Message: %d\n", sequence, message))
	}

	if sequence%Mod == 0 {
		//fmt.Printf("Consumer:: Sequence: %d, Message: %d\n", sequence, message)
	}
}
<commit_msg>ARM debugging.<commit_after>package main

import (
	"fmt"
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(3)

	producerSequence := NewSequence()
	consumerSequence1 := NewSequence()
	consumerSequence2 := NewSequence()

	producerBarrier := NewBarrier(producerSequence)
	consumerBarrier := NewBarrier(consumerSequence1, consumerSequence2)

	sequencer := NewSingleProducerSequencer(producerSequence, RingSize, consumerBarrier)
	go consume(producerBarrier, producerSequence, consumerSequence1)
	go consume(producerBarrier, producerSequence, consumerSequence2)

	started := time.Now()
	for i := int64(0); i < 10; i++ {
		ticket := sequencer.Next(1)
		ringBuffer[ticket&RingMask] = ticket
		sequencer.Publish(ticket)
		if ticket%Mod == 0 && ticket > 0 {
			finished := time.Now()
			elapsed := finished.Sub(started)
			fmt.Println(ticket, elapsed)
			started = time.Now()
		}
	}

	time.Sleep(time.Millisecond * 100)
}

func consume(barrier Barrier, source, sequence *Sequence) {
	worker := NewWorker(barrier, TestHandler{}, source, sequence)

	for {
		worker.Process()
	}
}

const Mod = 1000000 * 100 // 1 million * 10
const RingSize = 2
const RingMask = RingSize - 1

var ringBuffer [RingSize]int64

type TestHandler struct{}

func (this TestHandler) Consume(sequence, remaining int64) {
	message := ringBuffer[sequence&RingMask]
	if message != sequence {
		fmt.Printf("ERROR Consumer:: Sequence: %d, Message: %d\n", sequence, message)
		panic(fmt.Sprintf("Consumer:: Sequence: %d, Message: %d\n", sequence, message))
	}

	if sequence%Mod == 0 {
		//fmt.Printf("Consumer:: Sequence: %d, Message: %d\n", sequence, message)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"mime"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	mgo "gopkg.in/mgo.v2"

	"github.com/ONSdigital/florence/assets"
	"github.com/ONSdigital/florence/healthcheck"
	"github.com/ONSdigital/florence/upload"
	"github.com/ONSdigital/go-ns/handlers/reverseProxy"
	hc "github.com/ONSdigital/go-ns/healthcheck"
	"github.com/ONSdigital/go-ns/log"
	"github.com/ONSdigital/go-ns/server"
	"github.com/gorilla/pat"
	"github.com/gorilla/websocket"
)

var bindAddr = ":8080"
var babbageURL = "http://localhost:8080"
var zebedeeURL = "http://localhost:8082"
var recipeAPIURL = "http://localhost:22300"
var importAPIURL = "http://localhost:21800"
var uploadBucketName = "dp-frontend-florence-file-uploads"
var enableNewApp = false
var mongoURI = "localhost:27017"

var getAsset = assets.Asset
var upgrader = websocket.Upgrader{}
var session *mgo.Session

// Version is set by the make target
var Version string

func main() {
	log.Debug("florence version", log.Data{"version": Version})

	if v := os.Getenv("BIND_ADDR"); len(v) > 0 {
		bindAddr = v
	}
	if v := os.Getenv("BABBAGE_URL"); len(v) > 0 {
		babbageURL = v
	}
	if v := os.Getenv("ZEBEDEE_URL"); len(v) > 0 {
		zebedeeURL = v
	}
	if v := os.Getenv("RECIPE_API_URL"); len(v) > 0 {
		recipeAPIURL = v
	}
	if v := os.Getenv("UPLOAD_BUCKET_NAME"); len(v) > 0 {
		uploadBucketName = v
	}
	if v := os.Getenv("IMPORT_API_URL"); len(v) > 0 {
		recipeAPIURL = v
	}
	if v := os.Getenv("ENABLE_NEW_APP"); len(v) > 0 {
		enableNewApp, _ = strconv.ParseBool(v)
	}

	log.Namespace = "florence"

	zc := healthcheck.New(zebedeeURL, "zebedee")
	bc := healthcheck.New(babbageURL, "babbage")
	rc := healthcheck.New(recipeAPIURL, "recipe-api")
	ic := healthcheck.New(importAPIURL, "import-api")

	/*
		NOTE:
		If there's any issues with this Florence server proxying redirects
		from either Babbage or Zebedee then the code in the previous Java
		Florence server might give some clues for a solution: https://github.com/ONSdigital/florence/blob/b13df0708b30493b98e9ce239103c59d7f409f98/src/main/java/com/github/onsdigital/florence/filter/Proxy.java#L125-L135

		The code has purposefully not been included in this Go replacement
		because we can't see what issue it's fixing and whether it's necessary.
	*/

	babbageURL, err := url.Parse(babbageURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	babbageProxy := reverseProxy.Create(babbageURL, nil)

	zebedeeURL, err := url.Parse(zebedeeURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	zebedeeProxy := reverseProxy.Create(zebedeeURL, zebedeeDirector)

	recipeAPIURL, err := url.Parse(recipeAPIURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	recipeAPIProxy := reverseProxy.Create(recipeAPIURL, nil)

	importAPIURL, err := url.Parse(importAPIURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	importAPIProxy := reverseProxy.Create(importAPIURL, importAPIDirectory)

	router := pat.New()

	newAppHandler := refactoredIndexFile

	if !enableNewApp {
		newAppHandler = legacyIndexFile
	}

	uploader, err := upload.New(uploadBucketName)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}

	router.Path("/healthcheck").HandlerFunc(hc.Do)

	router.Path("/upload").Methods("GET").HandlerFunc(uploader.CheckUploaded)
	router.Path("/upload").Methods("POST").HandlerFunc(uploader.Upload)
	router.Path("/upload/{id}").Methods("GET").HandlerFunc(uploader.GetS3URL)

	router.Handle("/zebedee{uri:/.*}", zebedeeProxy)
	router.Handle("/recipes{uri:.*}", recipeAPIProxy)
	router.Handle("/import{uri:.*}", importAPIProxy)
	router.HandleFunc("/florence/dist/{uri:.*}", staticFiles)
	router.HandleFunc("/florence", legacyIndexFile)
	router.HandleFunc("/florence/", redirectToFlorence)
	router.HandleFunc("/florence/index.html", redirectToFlorence)
	router.HandleFunc("/florence/collections", legacyIndexFile)
	router.HandleFunc("/florence/publishing-queue", legacyIndexFile)
	router.HandleFunc("/florence/reports", legacyIndexFile)
	router.HandleFunc("/florence/users-and-access", legacyIndexFile)
	router.HandleFunc("/florence/websocket", websocketHandler)
	router.HandleFunc("/florence{uri:/.*}", newAppHandler)
	router.Handle("/{uri:.*}", babbageProxy)

	log.Debug("Starting server", log.Data{
		"bind_addr":      bindAddr,
		"babbage_url":    babbageURL,
		"zebedee_url":    zebedeeURL,
		"recipe_api_url": recipeAPIURL,
		"import_api_url": importAPIURL,
		"enable_new_app": enableNewApp,
	})

	s := server.New(bindAddr, router)
	// TODO need to reconsider default go-ns server timeouts
	s.Server.IdleTimeout = 120 * time.Second
	s.Server.WriteTimeout = 120 * time.Second
	s.Server.ReadTimeout = 30 * time.Second
	s.HandleOSSignals = false
	s.MiddlewareOrder = []string{"RequestID", "Log"}

	// FIXME temporary hack to remove timeout middleware (doesn't support hijacker interface)
	mo := s.MiddlewareOrder
	var newMo []string
	for _, mw := range mo {
		if mw != "Timeout" {
			newMo = append(newMo, mw)
		}
	}
	s.MiddlewareOrder = newMo

	go func() {
		if err := s.ListenAndServe(); err != nil {
			log.Error(err, nil)
			os.Exit(2)
		}
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, os.Kill)

	for {
		hc.MonitorExternal(bc, zc, ic, rc)

		timer := time.NewTimer(time.Second * 60)

		select {
		case <-timer.C:
			continue
		case <-stop:
			log.Info("shutting service down gracefully", nil)
			timer.Stop()
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			if err := s.Server.Shutdown(ctx); err != nil {
				log.Error(err, nil)
			}
			return
		}
	}
}

func redirectToFlorence(w http.ResponseWriter, req *http.Request) {
	http.Redirect(w, req, "/florence", 301)
}

func staticFiles(w http.ResponseWriter, req *http.Request) {
	path := req.URL.Query().Get(":uri")

	b, err := getAsset("../dist/" + path)
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, mime.TypeByExtension(filepath.Ext(path)))
	w.WriteHeader(200)
	w.Write(b)
}

func legacyIndexFile(w http.ResponseWriter, req *http.Request) {
	log.Debug("Getting legacy HTML file", nil)

	b, err := getAsset("../dist/legacy-assets/index.html")
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, "text/html")
	w.WriteHeader(200)
	w.Write(b)
}

func refactoredIndexFile(w http.ResponseWriter, req *http.Request) {
	log.Debug("Getting refactored HTML file", nil)

	b, err := getAsset("../dist/refactored.html")
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, "text/html")
	w.WriteHeader(200)
	w.Write(b)
}

func zebedeeDirector(req *http.Request) {
	if c, err := req.Cookie(`access_token`); err == nil && len(c.Value) > 0 {
		req.Header.Set(`X-Florence-Token`, c.Value)
	}
	req.URL.Path = strings.TrimPrefix(req.URL.Path, "/zebedee")
}

func importAPIDirectory(req *http.Request) {
	req.URL.Path = strings.TrimPrefix(req.URL.Path, "/import")
}

func websocketHandler(w http.ResponseWriter, req *http.Request) {
	c, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		log.ErrorR(req, err, nil)
		return
	}

	defer c.Close()

	err = c.WriteJSON(florenceServerEvent{"version", florenceVersionPayload{Version: Version}})
	if err != nil {
		log.ErrorR(req, err, nil)
		return
	}

	for {
		_, message, err := c.ReadMessage()
		if err != nil {
			log.ErrorR(req, err, nil)
			break
		}

		rdr := bufio.NewReader(bytes.NewReader(message))
		b, err := rdr.ReadBytes('{')
		if err != nil {
			log.ErrorR(req, err, log.Data{"bytes": string(b)})
			continue
		}

		tags := strings.Split(string(b), ":")
		eventID := tags[0]
		eventType := tags[1]
		eventData := message[len(eventID)+len(eventType)+2:]

		switch eventType {
		case "log":
			var e florenceLogEvent
			e.ServerTimestamp = time.Now().UTC().Format("2006-01-02T15:04:05.000-0700Z")
			err = json.Unmarshal(eventData, &e)
			if err != nil {
				log.ErrorR(req, err, log.Data{"data": string(eventData)})
				continue
			}
			log.Debug("client log", log.Data{"data": e})

			err = c.WriteJSON(florenceServerEvent{"ack", eventID})
			if err != nil {
				log.ErrorR(req, err, nil)
			}
		default:
			log.DebugR(req, "unknown event type", log.Data{"type": eventType, "data": string(eventData)})
		}

		// err = c.WriteMessage(mt, message)
		// if err != nil {
		// 	log.ErrorR(req, err, nil)
		// 	break
		// }
	}
}

type florenceLogEvent struct {
	ServerTimestamp string      `json:"-"`
	ClientTimestamp time.Time   `json:"clientTimestamp"`
	Type            string      `json:"type"`
	Location        string      `json:"location"`
	InstanceID      int         `json:"instanceID"`
	Payload         interface{} `json:"payload"`
}

type florenceServerEvent struct {
	Type    string      `json:"type"`
	Payload interface{} `json:"payload"`
}

type florenceVersionPayload struct {
	Version string `json:"version"`
}
<commit_msg>Fix var misassignment for import API<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"mime"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	mgo "gopkg.in/mgo.v2"

	"github.com/ONSdigital/florence/assets"
	"github.com/ONSdigital/florence/healthcheck"
	"github.com/ONSdigital/florence/upload"
	"github.com/ONSdigital/go-ns/handlers/reverseProxy"
	hc "github.com/ONSdigital/go-ns/healthcheck"
	"github.com/ONSdigital/go-ns/log"
	"github.com/ONSdigital/go-ns/server"
	"github.com/gorilla/pat"
	"github.com/gorilla/websocket"
)

var bindAddr = ":8080"
var babbageURL = "http://localhost:8080"
var zebedeeURL = "http://localhost:8082"
var recipeAPIURL = "http://localhost:22300"
var importAPIURL = "http://localhost:21800"
var uploadBucketName = "dp-frontend-florence-file-uploads"
var enableNewApp = false
var mongoURI = "localhost:27017"

var getAsset = assets.Asset
var upgrader = websocket.Upgrader{}
var session *mgo.Session

// Version is set by the make target
var Version string

func main() {
	log.Debug("florence version", log.Data{"version": Version})

	if v := os.Getenv("BIND_ADDR"); len(v) > 0 {
		bindAddr = v
	}
	if v := os.Getenv("BABBAGE_URL"); len(v) > 0 {
		babbageURL = v
	}
	if v := os.Getenv("ZEBEDEE_URL"); len(v) > 0 {
		zebedeeURL = v
	}
	if v := os.Getenv("RECIPE_API_URL"); len(v) > 0 {
		recipeAPIURL = v
	}
	if v := os.Getenv("UPLOAD_BUCKET_NAME"); len(v) > 0 {
		uploadBucketName = v
	}
	if v := os.Getenv("IMPORT_API_URL"); len(v) > 0 {
		importAPIURL = v
	}
	if v := os.Getenv("ENABLE_NEW_APP"); len(v) > 0 {
		enableNewApp, _ = strconv.ParseBool(v)
	}

	log.Namespace = "florence"

	zc := healthcheck.New(zebedeeURL, "zebedee")
	bc := healthcheck.New(babbageURL, "babbage")
	rc := healthcheck.New(recipeAPIURL, "recipe-api")
	ic := healthcheck.New(importAPIURL, "import-api")

	/*
		NOTE:
		If there's any issues with this Florence server proxying redirects
		from either Babbage or Zebedee then the code in the previous Java
		Florence server might give some clues for a solution: https://github.com/ONSdigital/florence/blob/b13df0708b30493b98e9ce239103c59d7f409f98/src/main/java/com/github/onsdigital/florence/filter/Proxy.java#L125-L135

		The code has purposefully not been included in this Go replacement
		because we can't see what issue it's fixing and whether it's necessary.
	*/

	babbageURL, err := url.Parse(babbageURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	babbageProxy := reverseProxy.Create(babbageURL, nil)

	zebedeeURL, err := url.Parse(zebedeeURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	zebedeeProxy := reverseProxy.Create(zebedeeURL, zebedeeDirector)

	recipeAPIURL, err := url.Parse(recipeAPIURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	recipeAPIProxy := reverseProxy.Create(recipeAPIURL, nil)

	importAPIURL, err := url.Parse(importAPIURL)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}
	importAPIProxy := reverseProxy.Create(importAPIURL, importAPIDirectory)

	router := pat.New()

	newAppHandler := refactoredIndexFile

	if !enableNewApp {
		newAppHandler = legacyIndexFile
	}

	uploader, err := upload.New(uploadBucketName)
	if err != nil {
		log.Error(err, nil)
		os.Exit(1)
	}

	router.Path("/healthcheck").HandlerFunc(hc.Do)

	router.Path("/upload").Methods("GET").HandlerFunc(uploader.CheckUploaded)
	router.Path("/upload").Methods("POST").HandlerFunc(uploader.Upload)
	router.Path("/upload/{id}").Methods("GET").HandlerFunc(uploader.GetS3URL)

	router.Handle("/zebedee{uri:/.*}", zebedeeProxy)
	router.Handle("/recipes{uri:.*}", recipeAPIProxy)
	router.Handle("/import{uri:.*}", importAPIProxy)
	router.HandleFunc("/florence/dist/{uri:.*}", staticFiles)
	router.HandleFunc("/florence", legacyIndexFile)
	router.HandleFunc("/florence/", redirectToFlorence)
	router.HandleFunc("/florence/index.html", redirectToFlorence)
	router.HandleFunc("/florence/collections", legacyIndexFile)
	router.HandleFunc("/florence/publishing-queue", legacyIndexFile)
	router.HandleFunc("/florence/reports", legacyIndexFile)
	router.HandleFunc("/florence/users-and-access", legacyIndexFile)
	router.HandleFunc("/florence/websocket", websocketHandler)
	router.HandleFunc("/florence{uri:/.*}", newAppHandler)
	router.Handle("/{uri:.*}", babbageProxy)

	log.Debug("Starting server", log.Data{
		"bind_addr":      bindAddr,
		"babbage_url":    babbageURL,
		"zebedee_url":    zebedeeURL,
		"recipe_api_url": recipeAPIURL,
		"import_api_url": importAPIURL,
		"enable_new_app": enableNewApp,
	})

	s := server.New(bindAddr, router)
	// TODO need to reconsider default go-ns server timeouts
	s.Server.IdleTimeout = 120 * time.Second
	s.Server.WriteTimeout = 120 * time.Second
	s.Server.ReadTimeout = 30 * time.Second
	s.HandleOSSignals = false
	s.MiddlewareOrder = []string{"RequestID", "Log"}

	// FIXME temporary hack to remove timeout middleware (doesn't support hijacker interface)
	mo := s.MiddlewareOrder
	var newMo []string
	for _, mw := range mo {
		if mw != "Timeout" {
			newMo = append(newMo, mw)
		}
	}
	s.MiddlewareOrder = newMo

	go func() {
		if err := s.ListenAndServe(); err != nil {
			log.Error(err, nil)
			os.Exit(2)
		}
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, os.Kill)

	for {
		hc.MonitorExternal(bc, zc, ic, rc)

		timer := time.NewTimer(time.Second * 60)

		select {
		case <-timer.C:
			continue
		case <-stop:
			log.Info("shutting service down gracefully", nil)
			timer.Stop()
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			if err := s.Server.Shutdown(ctx); err != nil {
				log.Error(err, nil)
			}
			return
		}
	}
}

func redirectToFlorence(w http.ResponseWriter, req *http.Request) {
	http.Redirect(w, req, "/florence", 301)
}

func staticFiles(w http.ResponseWriter, req *http.Request) {
	path := req.URL.Query().Get(":uri")

	b, err := getAsset("../dist/" + path)
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, mime.TypeByExtension(filepath.Ext(path)))
	w.WriteHeader(200)
	w.Write(b)
}

func legacyIndexFile(w http.ResponseWriter, req *http.Request) {
	log.Debug("Getting legacy HTML file", nil)

	b, err := getAsset("../dist/legacy-assets/index.html")
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, "text/html")
	w.WriteHeader(200)
	w.Write(b)
}

func refactoredIndexFile(w http.ResponseWriter, req *http.Request) {
	log.Debug("Getting refactored HTML file", nil)

	b, err := getAsset("../dist/refactored.html")
	if err != nil {
		log.Error(err, nil)
		w.WriteHeader(404)
		return
	}

	w.Header().Set(`Content-Type`, "text/html")
	w.WriteHeader(200)
	w.Write(b)
}

func zebedeeDirector(req *http.Request) {
	if c, err := req.Cookie(`access_token`); err == nil && len(c.Value) > 0 {
		req.Header.Set(`X-Florence-Token`, c.Value)
	}
	req.URL.Path = strings.TrimPrefix(req.URL.Path, "/zebedee")
}

func importAPIDirectory(req *http.Request) {
	req.URL.Path = strings.TrimPrefix(req.URL.Path, "/import")
}

func websocketHandler(w http.ResponseWriter, req *http.Request) {
	c, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		log.ErrorR(req, err, nil)
		return
	}

	defer c.Close()

	err = c.WriteJSON(florenceServerEvent{"version", florenceVersionPayload{Version: Version}})
	if err != nil {
		log.ErrorR(req, err, nil)
		return
	}

	for {
		_, message, err := c.ReadMessage()
		if err != nil {
			log.ErrorR(req, err, nil)
			break
		}

		rdr := bufio.NewReader(bytes.NewReader(message))
		b, err := rdr.ReadBytes('{')
		if err != nil {
			log.ErrorR(req, err, log.Data{"bytes": string(b)})
			continue
		}

		tags := strings.Split(string(b), ":")
		eventID := tags[0]
		eventType := tags[1]
		eventData := message[len(eventID)+len(eventType)+2:]

		switch eventType {
		case "log":
			var e florenceLogEvent
			e.ServerTimestamp = time.Now().UTC().Format("2006-01-02T15:04:05.000-0700Z")
			err = json.Unmarshal(eventData, &e)
			if err != nil {
				log.ErrorR(req, err, log.Data{"data": string(eventData)})
				continue
			}
			log.Debug("client log", log.Data{"data": e})

			err = c.WriteJSON(florenceServerEvent{"ack", eventID})
			if err != nil {
				log.ErrorR(req, err, nil)
			}
		default:
			log.DebugR(req, "unknown event type", log.Data{"type": eventType, "data": string(eventData)})
		}

		// err = c.WriteMessage(mt, message)
		// if err != nil {
		// 	log.ErrorR(req, err, nil)
		// 	break
		// }
	}
}

type florenceLogEvent struct {
	ServerTimestamp string      `json:"-"`
	ClientTimestamp time.Time   `json:"clientTimestamp"`
	Type            string      `json:"type"`
	Location        string      `json:"location"`
	InstanceID      int         `json:"instanceID"`
	Payload         interface{} `json:"payload"`
}

type florenceServerEvent struct {
	Type    string      `json:"type"`
	Payload interface{} `json:"payload"`
}

type florenceVersionPayload struct {
	Version string `json:"version"`
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"log/syslog"
	"net/smtp"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
	_ "github.com/go-sql-driver/mysql"
	"github.com/johansundell/cocapi"
)

var db *sql.DB
var mysqlUser, mysqlPass, mysqlDb, mysqlHost string
var queryInsertUpdateMember = `INSERT INTO members (tag, name, created, last_updated, active) VALUES (?, ?, null, null, 1) ON DUPLICATE KEY UPDATE member_id=LAST_INSERT_ID(member_id), last_updated = NOW(), active = 1`
var consumerKey, consumerSecret, accessToken, accessSecret string

func init() {

	mysqlDb = "cocsniffer"
	mysqlHost = os.Getenv("MYSQL_COC_HOST")
	mysqlUser = os.Getenv("MYSQL_USER")
	mysqlPass = os.Getenv("MYSQL_PASS")

	consumerKey = os.Getenv("TWITTER_CONSKEY")
	consumerSecret = os.Getenv("TWITTER_CONSSEC")
	accessToken = os.Getenv("TWITTER_ACCTOK")
	accessSecret = os.Getenv("TWITTER_ACCSEC")
}

func main() {
	useSyslog := flag.Bool("syslog", false, "Use syslog")
	flag.Parse()
	if *useSyslog {
		logwriter, e := syslog.New(syslog.LOG_NOTICE, "cocsniffer")
		if e == nil {
			log.SetOutput(logwriter)
		}
	}
	db, _ = sql.Open("mysql", mysqlUser+":"+mysqlPass+"@tcp("+mysqlHost+":3306)/"+mysqlDb)
	defer db.Close()

	getMembersData()
	ticker := time.NewTicker(5 * time.Minute)
	quit := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				getMembersData()
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	config := oauth1.NewConfig(consumerKey, consumerSecret)
	token := oauth1.NewToken(accessToken, accessSecret)
	// OAuth1 http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter client
	client := twitter.NewClient(httpClient)

	// Convenience Demux demultiplexed stream messages
	demux := twitter.NewSwitchDemux()
	demux.Tweet = func(tweet *twitter.Tweet) {
		//fmt.Println("found one", tweet.Text)
		if strings.Contains(strings.ToLower(tweet.Text), strings.ToLower("Maintenance break")) {
			sendEmail("johan@pixpro.net", "johan@sundell.com", "COC alert", tweet.Text)
			log.Println("Email sent:", tweet.Text)
		}
	}
	demux.DM = func(dm *twitter.DirectMessage) {
		//fmt.Println(dm.SenderID)
	}
	demux.Event = func(event *twitter.Event) {
		//fmt.Printf("%#v\n", event)
	}

	log.Println("Starting Stream...")

	// FILTER
	filterParams := &twitter.StreamFilterParams{
		Follow:        []string{"730400376", "240359880"},
		Track:         []string{"Maintenance", "Maintenance."},
		StallWarnings: twitter.Bool(true),
	}

	stream, err := client.Streams.Filter(filterParams)
	if err != nil {
		log.Fatal(err)
	}

	// Receive messages until stopped or stream quits
	go demux.HandleChan(stream.Messages)

	// Wait for SIGINT and SIGTERM (HIT CTRL-C)
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	close(quit)
	stream.Stop()
	fmt.Println("Bye ;)")
}

func getMembersData() {
	members, err := cocapi.GetMemberInfo()
	if err != nil {
		reportError(err)
	}

	var ids = make([]string, 0)
	for _, m := range members.Items {
		if result, err := db.Exec(queryInsertUpdateMember, m.Tag, m.Name); err != nil {
			log.Println(err)
		} else {
			if id, err := result.LastInsertId(); err != nil {
				log.Println(err)
			} else {
				ids = append(ids, strconv.Itoa(int(id)))
			}
		}
	}
	db.Exec("UPDATE members SET active = 0 WHERE member_id NOT IN (" + strings.Join(ids, ", ") + ")")
	log.Println("done members func")
}

func reportError(err error) {
	log.Println("Fatal error:", err)
	os.Exit(0)
}

func sendEmail(to, from, subject, message string) bool {
	body := "To: " + to + "\r\nSubject: " + subject + "\r\n\r\n" + message
	if err := smtp.SendMail("127.0.0.1:25", nil, from, []string{to}, []byte(body)); err != nil {
		log.Println(err)
		return false
	}
	return true
}
<commit_msg>Fixes<commit_after>package main

import (
	"database/sql"
	"flag"
	"log"
	"log/syslog"
	"net/smtp"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
	_ "github.com/go-sql-driver/mysql"
	"github.com/johansundell/cocapi"
)

var db *sql.DB
var mysqlUser, mysqlPass, mysqlDb, mysqlHost string
var queryInsertUpdateMember = `INSERT INTO members (tag, name, created, last_updated, active) VALUES (?, ?, null, null, 1) ON DUPLICATE KEY UPDATE member_id=LAST_INSERT_ID(member_id), last_updated = NOW(), active = 1`
var consumerKey, consumerSecret, accessToken, accessSecret string

func init() {

	mysqlDb = "cocsniffer"
	mysqlHost = os.Getenv("MYSQL_COC_HOST")
	mysqlUser = os.Getenv("MYSQL_USER")
	mysqlPass = os.Getenv("MYSQL_PASS")

	consumerKey = os.Getenv("TWITTER_CONSKEY")
	consumerSecret = os.Getenv("TWITTER_CONSSEC")
	accessToken = os.Getenv("TWITTER_ACCTOK")
	accessSecret = os.Getenv("TWITTER_ACCSEC")
}

func main() {
	useSyslog := flag.Bool("syslog", false, "Use syslog")
	flag.Parse()
	if *useSyslog {
		logwriter, e := syslog.New(syslog.LOG_NOTICE, "cocsniffer")
		if e == nil {
			log.SetOutput(logwriter)
		}
	}
	db, _ = sql.Open("mysql", mysqlUser+":"+mysqlPass+"@tcp("+mysqlHost+":3306)/"+mysqlDb)
	defer db.Close()

	getMembersData()
	ticker := time.NewTicker(5 * time.Minute)
	quit := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				getMembersData()
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	config := oauth1.NewConfig(consumerKey, consumerSecret)
	token := oauth1.NewToken(accessToken, accessSecret)
	// OAuth1 http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter client
	client := twitter.NewClient(httpClient)

	// Convenience Demux demultiplexed stream messages
	demux := twitter.NewSwitchDemux()
	demux.Tweet = func(tweet *twitter.Tweet) {
		//fmt.Println("found one", tweet.Text)
		if strings.Contains(strings.ToLower(tweet.Text), strings.ToLower("Maintenance break")) {
			sendEmail("johan@pixpro.net", "johan@sundell.com", "COC alert", tweet.Text)
			log.Println("Email sent:", tweet.Text)
		}
	}
	demux.DM = func(dm *twitter.DirectMessage) {
		//fmt.Println(dm.SenderID)
	}
	demux.Event = func(event *twitter.Event) {
		//fmt.Printf("%#v\n", event)
	}

	log.Println("Starting Stream...")

	// FILTER
	filterParams := &twitter.StreamFilterParams{
		Follow:        []string{"730400376", "240359880"},
		Track:         []string{"Maintenance", "Maintenance."},
		StallWarnings: twitter.Bool(true),
	}

	stream, err := client.Streams.Filter(filterParams)
	if err != nil {
		log.Fatal(err)
	}

	// Receive messages until stopped or stream quits
	go demux.HandleChan(stream.Messages)

	// Wait for SIGINT and SIGTERM (HIT CTRL-C)
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	close(quit)
	stream.Stop()
	log.Println("Bye ;)")
}

func getMembersData() {
	members, err := cocapi.GetMemberInfo()
	if err != nil {
		reportError(err)
	}

	var ids = make([]string, 0)
	for _, m := range members.Items {
		if result, err := db.Exec(queryInsertUpdateMember, m.Tag, m.Name); err != nil {
			log.Println(err)
		} else {
			if id, err := result.LastInsertId(); err != nil {
				log.Println(err)
			} else {
				ids = append(ids, strconv.Itoa(int(id)))
			}
		}
	}
	db.Exec("UPDATE members SET active = 0 WHERE member_id NOT IN (" + strings.Join(ids, ", ") + ")")
	log.Println("done members func")
}

func reportError(err error) {
	log.Println("Fatal error:", err)
	os.Exit(0)
}

func sendEmail(to, from, subject, message string) bool {
	body := "To: " + to + "\r\nSubject: " + subject + "\r\n\r\n" + message
	if err := smtp.SendMail("127.0.0.1:25", nil, from, []string{to}, []byte(body)); err != nil {
		log.Println(err)
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
)

type database struct {
	AppServer string
	DbServer  string
	DbName    string
	User      string
	Pass      string
}

var help = flag.Bool("help", false, "shows usage")

func init() {
	flag.BoolVar(help, "h", false, "shows usage")
}

func main() {
	flag.Parse()
	if *help {
		usage("")
	}

	databases := mustReadDatabasesConfigFile()

	if len(os.Args[1:]) == 0 {
		usage("Target database unspecified; where should I run the query?")
	}

	sql := readInput(os.Stdin)
	if len(sql) <= 3 {
		usage("No SQL to run. Exiting.")
	}

	targetDatabases := []string{}
	for _, k := range os.Args[1:] {
		if _, ok := databases[k]; k != "all" && !ok {
			usage("Target database unknown: [%v]", k)
		}
		if k == "all" {
			targetDatabases = nil
			for k := range databases {
				targetDatabases = append(targetDatabases, k)
			}
			break
		}
		targetDatabases = append(targetDatabases, k)
	}

	out := make(chan string)
	go println(out)

	quitContext, cancel := context.WithCancel(context.Background())
	go awaitSignal(cancel)

	var wg sync.WaitGroup
	wg.Add(len(targetDatabases))

	returnCode := 0
	for _, k := range targetDatabases {
		go func(db database, k string) {
			defer wg.Done()
			if r := runSQL(quitContext, db, sql, k, len(targetDatabases) > 1, out); !r {
				returnCode = 1
			}
		}(databases[k], k)
	}

	wg.Wait()
	os.Exit(returnCode)
}

func runSQL(quitContext context.Context, db database, sql string, key string, prependKey bool, out chan string) bool {
	userOption := ""
	if db.User != "" {
		userOption = fmt.Sprintf("-u %v ", db.User)
	}

	passOption := ""
	if db.Pass != "" {
		passOption = fmt.Sprintf("-p%v ", db.Pass)
	}

	hostOption := ""
	if db.DbServer != "" {
		hostOption = fmt.Sprintf("-h %v ", db.DbServer)
	}

	prepend := ""
	if prependKey {
		prepend = key + "\t"
	}

	mysql := "mysql"
	options := fmt.Sprintf(" -Nsr %v%v%v%v -e ", userOption, passOption, hostOption, db.DbName)

	var cmd *exec.Cmd
	if db.AppServer != "" {
		query := fmt.Sprintf(`'%v'`, strings.Replace(sql, `'`, `'"'"'`, -1))
		cmd = exec.CommandContext(quitContext, "ssh", db.AppServer, mysql+options+query)
	} else {
		args := append(trimEmpty(strings.Split(options, " ")), sql)
		cmd = exec.CommandContext(quitContext, "mysql", args...)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("Cannot create pipe for STDOUT of running command on %v; not running. err=%v\n", key, err)
		return false
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("Cannot create pipe for STDERR of running command on %v; not running. err=%v\n", key, err)
		return false
	}

	if err := cmd.Start(); err != nil {
		log.Printf("Cannot start command on %v; not running. err=%v\n", key, err)
		return false
	}

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		out <- prepend + scanner.Text()
	}

	stderrLines := []string{}
	scanner = bufio.NewScanner(stderr)
	for scanner.Scan() {
		stderrLines = append(stderrLines, scanner.Text())
	}

	cmd.Wait()

	result := true
	if len(stderrLines) > 0 {
		result = false
		log.Println(key + " had errors:")
		for _, v := range stderrLines {
			log.Println(key + " [ERROR] " + v)
		}
	}

	return result
}

func println(ss chan string) {
	for s := range ss {
		fmt.Println(s)
	}
}

func readInput(r io.Reader) string {
	ls := []string{}
	var err error
	rd := bufio.NewReader(r)

	for {
		var s string
		s, err = rd.ReadString('\n')
		if err == io.EOF {
			return strings.Join(ls, " ")
		}
		s = strings.TrimSpace(s)
		if len(s) == 0 {
			continue
		}
		ls = append(ls, strings.TrimSpace(s))
	}
}

func trimEmpty(s []string) []string {
	var r []string
	for _, str := range s {
		if str != "" {
			r = append(r, str)
		}
	}
	return r
}

func awaitSignal(cancel context.CancelFunc) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	<-signals
	cancel()
}
<commit_msg>remove unnecessary channel buffer<commit_after>package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
)

type database struct {
	AppServer string
	DbServer  string
	DbName    string
	User      string
	Pass      string
}

var help = flag.Bool("help", false, "shows usage")

func init() {
	flag.BoolVar(help, "h", false, "shows usage")
}

func main() {
	flag.Parse()
	if *help {
		usage("")
	}

	databases := mustReadDatabasesConfigFile()

	if len(os.Args[1:]) == 0 {
		usage("Target database unspecified; where should I run the query?")
	}

	sql := readInput(os.Stdin)
	if len(sql) <= 3 {
		usage("No SQL to run. Exiting.")
	}

	targetDatabases := []string{}
	for _, k := range os.Args[1:] {
		if _, ok := databases[k]; k != "all" && !ok {
			usage("Target database unknown: [%v]", k)
		}
		if k == "all" {
			targetDatabases = nil
			for k := range databases {
				targetDatabases = append(targetDatabases, k)
			}
			break
		}
		targetDatabases = append(targetDatabases, k)
	}

	out := make(chan string)
	go println(out)

	quitContext, cancel := context.WithCancel(context.Background())
	go awaitSignal(cancel)

	var wg sync.WaitGroup
	wg.Add(len(targetDatabases))

	returnCode := 0
	for _, k := range targetDatabases {
		go func(db database, k string) {
			defer wg.Done()
			if r := runSQL(quitContext, db, sql, k, len(targetDatabases) > 1, out); !r {
				returnCode = 1
			}
		}(databases[k], k)
	}

	wg.Wait()
	os.Exit(returnCode)
}

func runSQL(quitContext context.Context, db database, sql string, key string, prependKey bool, out chan string) bool {
	userOption := ""
	if db.User != "" {
		userOption = fmt.Sprintf("-u %v ", db.User)
	}

	passOption := ""
	if db.Pass != "" {
		passOption = fmt.Sprintf("-p%v ", db.Pass)
	}

	hostOption := ""
	if db.DbServer != "" {
		hostOption = fmt.Sprintf("-h %v ", db.DbServer)
	}

	prepend := ""
	if prependKey {
		prepend = key + "\t"
	}

	mysql := "mysql"
	options := fmt.Sprintf(" -Nsr %v%v%v%v -e ", userOption, passOption, hostOption, db.DbName)

	var cmd *exec.Cmd
	if db.AppServer != "" {
		query := fmt.Sprintf(`'%v'`, strings.Replace(sql, `'`, `'"'"'`, -1))
		cmd = exec.CommandContext(quitContext, "ssh", db.AppServer, mysql+options+query)
	} else {
		args := append(trimEmpty(strings.Split(options, " ")), sql)
		cmd = exec.CommandContext(quitContext, "mysql", args...)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("Cannot create pipe for STDOUT of running command on %v; not running. err=%v\n", key, err)
		return false
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("Cannot create pipe for STDERR of running command on %v; not running. err=%v\n", key, err)
		return false
	}

	if err := cmd.Start(); err != nil {
		log.Printf("Cannot start command on %v; not running. err=%v\n", key, err)
		return false
	}

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		out <- prepend + scanner.Text()
	}

	stderrLines := []string{}
	scanner = bufio.NewScanner(stderr)
	for scanner.Scan() {
		stderrLines = append(stderrLines, scanner.Text())
	}

	cmd.Wait()

	result := true
	if len(stderrLines) > 0 {
		result = false
		log.Println(key + " had errors:")
		for _, v := range stderrLines {
			log.Println(key + " [ERROR] " + v)
		}
	}

	return result
}

func println(ss chan string) {
	for s := range ss {
		fmt.Println(s)
	}
}

func readInput(r io.Reader) string {
	ls := []string{}
	var err error
	rd := bufio.NewReader(r)

	for {
		var s string
		s, err = rd.ReadString('\n')
		if err == io.EOF {
			return strings.Join(ls, " ")
		}
		s = strings.TrimSpace(s)
		if len(s) == 0 {
			continue
		}
		ls = append(ls, strings.TrimSpace(s))
	}
}

func trimEmpty(s []string) []string {
	var r []string
	for _, str := range s {
		if str != "" {
			r = append(r, str)
		}
	}
	return r
}

func awaitSignal(cancel context.CancelFunc) {
	signals := make(chan os.Signal)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	<-signals
	cancel()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

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

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

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

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

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

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

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

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

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

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

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

	// websocket endpoint
	v1Endpoints.GET("/websocket",
		v1.Websocket(
			connsHub.Len,
			memoryCache.GetLatestConfig,
			memoryCache.GetLatestIncomes,
			connsHub.Broadcast,
			hub.WrapPutWebsocketConn(connsHub.PutConn)),
	)

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

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

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

	return nil
}

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

func initStorage() {
	// storage service
	s, err := mysql.New(config.DB.DataSourceName)
	panicIfErrored(err)
	s.SetMaxOpenConns(config.DB.MaxOpenConns)
	s.SetMaxIdleConns(config.DB.MaxIdleConns)
	store = s
}

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

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

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

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

	// update bitcoin price on start
	updateBitcoinPrice()
}

func initHub() {
	connsHub = list.New()
}

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

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

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

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

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

	// broadcast bitcoin price to all users
	msg, _ := json.Marshal(models.WebsocketMessage{
		BitcoinPrice: utils.HumanReadableUSD(p),
	})
	connsHub.Broadcast(msg)

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

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

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

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

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

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

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

// sync cache with storage
func syncCache() {
	// update config in cache
	config, err := store.GetLatestConfig()
	if err != nil {
		fmt.Fprintf(panicWriter, "Update latest config error: %v", err)
		return
	}
	memoryCache.SetLatestConfig(config)

	// update rates in cache
	lessRates, err := store.GetRewardRatesByType(models.RewardRateTypeLess)
	if err != nil {
		fmt.Fprintf(panicWriter, "Update less rate error: %v", err)
		return
	}

	moreRates, err := store.GetRewardRatesByType(models.RewardRateTypeMore)
	if err != nil {
		fmt.Fprintf(panicWriter, "Update more rate error: %v", err)
		return
	}

	memoryCache.SetRewardRates(models.RewardRateTypeMore, moreRates)
	memoryCache.SetRewardRates(models.RewardRateTypeLess, lessRates)

	fmt.Fprintf(logWriter, "Successfully sync cache\n")
}

// fail fast on initialization
func panicIfErrored(err error) {
	if err != nil {
		// Tricky:
		// pass a nil *errors.Error into this function
		// err is not nil
		// see discussion here:
		// https://github.com/go-playground/validator/issues/134
		// or
		// http://stackoverflow.com/questions/29138591/hiding-nil-values-understanding-why-golang-fails-here/29138676#29138676
		if e, ok := err.(*errors.Error); ok {
			if e != nil {
				panic(e.ErrStringForLogging)
			}
		} else {
			panic(err)
		}
	}
}
<commit_msg>Do not recover from create withdrawals panic as it is critical<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

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

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

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

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

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

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

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

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

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

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

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

	// websocket endpoint
	v1Endpoints.GET("/websocket",
		v1.Websocket(
			connsHub.Len,
			memoryCache.GetLatestConfig,
			memoryCache.GetLatestIncomes,
			connsHub.Broadcast,
			hub.WrapPutWebsocketConn(connsHub.PutConn)),
	)

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

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

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

	return nil
}

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

func initStorage() {
	// storage service
	s, err := mysql.New(config.DB.DataSourceName)
	panicIfErrored(err)
	s.SetMaxOpenConns(config.DB.MaxOpenConns)
	s.SetMaxIdleConns(config.DB.MaxIdleConns)
	store = s
}

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

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

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

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

	// update bitcoin price on start
	updateBitcoinPrice()
}

func initHub() {
	connsHub = list.New()
}

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

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

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

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

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

	// broadcast bitcoin price to all users
	msg, _ := json.Marshal(models.WebsocketMessage{
		BitcoinPrice: utils.HumanReadableUSD(p),
	})
	connsHub.Broadcast(msg)

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

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

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

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

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

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

// sync cache with storage
func syncCache() {
	// update config in cache
	config, err := store.GetLatestConfig()
	if err != nil {
		fmt.Fprintf(panicWriter, "Update latest config error: %v", err)
		return
	}
	memoryCache.SetLatestConfig(config)

	// update rates in cache
	lessRates, err := store.GetRewardRatesByType(models.RewardRateTypeLess)
	if err != nil {
		fmt.Fprintf(panicWriter, "Update less rate error: %v", err)
		return
	}

	moreRates, err := store.GetRewardRatesByType(models.RewardRateTypeMore)
	if err != nil {
		fmt.Fprintf(panicWriter, "Update more rate error: %v", err)
		return
	}

	memoryCache.SetRewardRates(models.RewardRateTypeMore, moreRates)
	memoryCache.SetRewardRates(models.RewardRateTypeLess, lessRates)

	fmt.Fprintf(logWriter, "Successfully sync cache\n")
}

// fail fast on initialization
func panicIfErrored(err error) {
	if err != nil {
		// Tricky:
		// pass a nil *errors.Error into this function
		// err is not nil
		// see discussion here:
		// https://github.com/go-playground/validator/issues/134
		// or
		// http://stackoverflow.com/questions/29138591/hiding-nil-values-understanding-why-golang-fails-here/29138676#29138676
		if e, ok := err.(*errors.Error); ok {
			if e != nil {
				panic(e.ErrStringForLogging)
			}
		} else {
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/alienth/fastlyctl/util"
	"github.com/alienth/go-fastly"
	"github.com/juju/ratelimit"
	"github.com/urfave/cli"
	"gopkg.in/mcuadros/go-syslog.v2"
)

// getIPs takes in an haproxy log line and returns the client IP and the CDN
// connecting IP which is in the captured request header. For us, the CDN
// connecting IP happens to be the 8th captured request header.
func getIPs(logLine string) (*net.IP, *net.IP) {
	if logLine == "" {
		return nil, nil
	}
	// This string parsing stuff was lifted from TPS
	var a, b int
	if a = strings.Index(logLine, "]:") + 3; a == -1 {
		return nil, nil
	}
	if b = strings.Index(logLine[a:], ":"); b == -1 {
		return nil, nil
	}
	clientIPString := logLine[a : a+b]
	clientIP := net.ParseIP(clientIPString)
	if clientIP == nil {
		return nil, nil
	}

	logLine = logLine[a+b:]
	// The first curly-braced block contains our request headers
	if a = strings.Index(logLine, "{"); a == -1 {
		return &clientIP, nil
	}
	if b = strings.Index(logLine[a:], "}"); b == -1 {
		return &clientIP, nil
	}
	bracketedHeaders := logLine[a : a+b]
	headers := strings.Split(bracketedHeaders, "|")
	if len(headers) < 7 {
		return &clientIP, nil
	}
	ipString := headers[7]
	cdnIP := net.ParseIP(ipString)
	if cdnIP == nil {
		return &clientIP, nil
	}
	return &clientIP, &cdnIP
}

func getLogTimestamp(logLine string) int64 {
	timestampStr := strings.Split(logLine, " ")[6]
	ts, _ := time.Parse("[02/Jan/2006:15:04:05.999]", timestampStr)

	return (ts.Unix())
}

type ipList struct {
	ips  []net.IP
	nets []net.IPNet
}

func readIPList(filename string) (*ipList, error) {
	var list ipList
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Index(strings.TrimSpace(line), "#") == 0 {
			continue
		}
		if strings.TrimSpace(line) == "" {
			continue
		}
		if strings.Contains(line, "/") {
			_, n, err := net.ParseCIDR(line)
			if err != nil {
				return nil, fmt.Errorf("Unable to parse CIDR.\nLine:\n%s\nError:\n%s\n", line, err)
			}
			if n != nil {
				list.nets = append(list.nets, *n)
			}
		} else {
			ip := net.ParseIP(line)
			if ip != nil {
				list.ips = append(list.ips, ip)
			} else {
				return nil, fmt.Errorf("Unable to parse IP address in list: %s\n", line)
			}
		}

	}
	return &list, nil
}

func (l *ipList) contains(checkIP *net.IP) bool {
	if l == nil {
		return false
	}
	for _, ip := range l.ips {
		if ip.Equal(*checkIP) {
			return true
		}
	}
	for _, net := range l.nets {
		if net.Contains(*checkIP) {
			return true
		}
	}

	return false
}

type ipRate struct {
	ip      *net.IP
	bucket  *ratelimit.Bucket
	entries []*util.ACLEntry

	FirstHit    int64 `json:"first_hit,omitempty"`
	LastHit     int64 `json:"last_hit,omitempty"`
	LastLimit   int64 `json:"last_limit,omitempty"`
	Hits        int   `json:"hits,omitempty"`
	Strikes     int   `json:"strikes,omitempty"`
	Expire      int64 `json:"-"`
	LimitExpire int64 `json:"limit_expire,omitempty"`
}

func (ipr *ipRate) New(ip *net.IP) {
	ipr.ip = ip
	ipr.bucket = ratelimit.NewBucket(time.Duration(3)*time.Minute, 180)
	ipr.Expire = time.Now().Add(time.Duration(1) * time.Hour).Unix()
}

// Records a hit and returns true if it is over limit.
func (ipr *ipRate) Hit() bool {
	// Here we pretend that the IP performed no hits before the limit was removed,
	// since in testing the IP is not actually limited.
	if ipr.LastLimit != 0 && time.Now().Sub(time.Unix(ipr.LastLimit, 0)) < time.Duration(1)*time.Hour {
		return false
	}

	_, isSoonerThanMaxWait := ipr.bucket.TakeMaxDuration(1, 0)
	if ipr.FirstHit == 0 {
		ipr.FirstHit = time.Now().Unix()
	}
	ipr.LastHit = time.Now().Unix()
	ipr.Hits++
	ipr.Expire = time.Now().Add(time.Duration(1) * time.Hour).Unix()
	return !isSoonerThanMaxWait
}

// Limit adds an IP to a fastly edge ACL
func (ipr *ipRate) Limit() error {

	service, err := util.GetServiceByNameOrID(client, "teststackoverflow.com")
	if err != nil {
		return err
	}

	acl, err := util.NewACL(client, service.ID, "ratelimit")
	if err != nil {
		return err
	}

	// TODO this is heavy.. cache? do we even need?(see below)
	entries, err := acl.ListEntries()
	if err != nil {
		return err
	}

	var entry *util.ACLEntry

	for _, e := range entries {
		if ipr.ip.String() == e.IP {
			fmt.Printf("IP %s is already limited.\n", e.IP)
			// Not all ACL entries were necessarily created by this script, so silently ignore unmarshal errors.
			if err := json.Unmarshal([]byte(e.Comment), ipr); err != nil {
				break
			}
			entry = &util.ACLEntry{Client: client, ID: e.ID, ACLID: e.ACLID, ServiceID: service.ID}
			break
		}
	}

	ipr.LastLimit = time.Now().Unix()
	ipr.Strikes++
	ipr.LimitExpire = time.Now().Add(time.Duration(5) * time.Minute).Unix()
	ipr.Expire = time.Now().Add(time.Duration(24) * time.Hour).Unix()
	comment, err := json.Marshal(ipr)
	if err != nil {
		return err
	}
	if entry == nil {
		entry, err = util.NewACLEntry(client, service.ID, "ratelimit", ipr.ip.String(), 0, string(comment), false)
		if err != nil {
			return err
		}
		entry.Add()
	} else {
		entry.Comment = string(comment)
		//		entry.Update()
	}

	ipr.entries = append(ipr.entries, entry)
	return nil
}

// Removes an IP from ratelimits
func (ipr *ipRate) RemoveLimit() error {
	if len(ipr.entries) > 0 {
		fmt.Printf("Unlimiting IP %s\n", ipr.ip.String())
		for i, entry := range ipr.entries {
			if err := entry.Remove(); err != nil {
				fmt.Println(err)
			}
			ipr.entries[i] = nil
		}
		ipr.entries = ipr.entries[:0]
	}
	return nil
}

type hitMap struct {
	sync.RWMutex
	m map[string]*ipRate
}

func (hits *hitMap) expireRecords() {
	for {
		hits.Lock()
		for ip, ipr := range hits.m {
			if ipr.Expire < time.Now().Unix() {
				ipr.RemoveLimit()
				delete(hits.m, ip)
			}
		}
		hits.Unlock()
		time.Sleep(time.Duration(60) * time.Second)
	}
}

func (hits *hitMap) expireLimits() {
	for {
		hits.Lock()
		for _, ipr := range hits.m {
			if ipr.LimitExpire < time.Now().Unix() {
				ipr.RemoveLimit()
			}
		}
		hits.Unlock()
		time.Sleep(time.Duration(15) * time.Second)
	}
}

// Fetches down remote ACLs and populates local hitMap with previously stored data.
func (hits *hitMap) importIPRates() error {
	service, err := util.GetServiceByNameOrID(client, "teststackoverflow.com")
	if err != nil {
		return err
	}

	acl, err := util.NewACL(client, service.ID, "ratelimit")
	if err != nil {
		return err
	}

	entries, err := acl.ListEntries()
	if err != nil {
		return err
	}

	// TODO: work with IPs existing in multiple services.
	for _, e := range entries {
		var ipr *ipRate
		var found bool
		hits.Lock()
		if ipr, found = hits.m[e.IP]; !found {
			ipr = &ipRate{}
			ip := net.ParseIP(e.IP)
			if ip == nil {
				return fmt.Errorf("Unable to parse IP %s in ACL.")
			}
			ipr.New(&ip)
			hits.m[ip.String()] = ipr
		}
		hits.Unlock()
		if err := json.Unmarshal([]byte(e.Comment), ipr); err != nil {
			break
		}
		entry := &util.ACLEntry{Client: client, ID: e.ID, ACLID: e.ACLID, ServiceID: service.ID}
		ipr.entries = append(ipr.entries, entry)
	}

	return nil
}

var client *fastly.Client

func main() {
	app := cli.NewApp()
	app.Name = "fastly-ratelimit"

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "whitelist, w",
			Usage: "Read `FILE` containing addresses which should whitelisted from any bans.",
		},
		cli.StringFlag{
			Name:  "listen, l",
			Usage: "Specify listen `ADDRESS:PORT`.",
			Value: "0.0.0.0:514",
		},
		cli.StringFlag{
			Name:   "fastly-key, K",
			Usage:  "Fastly API Key. Can be read from 'fastly_key' file in CWD.",
			EnvVar: "FASTLY_KEY",
			Value:  util.GetFastlyKey(),
		},
	}
	app.Action = func(c *cli.Context) error {
		client, _ = fastly.NewClient(c.GlobalString("fastly-key"))
		channel := make(syslog.LogPartsChannel)
		handler := syslog.NewChannelHandler(channel)

		server := syslog.NewServer()
		server.SetFormat(syslog.RFC3164)
		server.SetHandler(handler)
		if err := server.ListenUDP(c.GlobalString("listen")); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to listen: %s\n", err), -1)
		}
		if err := server.Boot(); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to start server: %s\n", err), -1)
		}

		var whitelist *ipList
		var err error
		whitelistFile := c.GlobalString("whitelist")
		if whitelistFile != "" {
			if whitelist, err = readIPList(whitelistFile); err != nil {
				return cli.NewExitError(fmt.Sprintf("Error reading whitelist file: %s", err), -1)
			}
		}

		var hits = hitMap{m: make(map[string]*ipRate)}
		if err := hits.importIPRates(); err != nil {
			return cli.NewExitError(fmt.Sprintf("Error importing existing IP rates: %s", err), -1)
		}
		go hits.expireRecords()
		go hits.expireLimits()
		go func(channel syslog.LogPartsChannel) {
			for logParts := range channel {
				var line string
				var ok bool
				if line, ok = logParts["content"].(string); !ok || line == "" {
					continue
				}
				clientIP, cdnIP := getIPs(line)
				if cdnIP == nil || clientIP == nil {
					continue
				}
				var ipr *ipRate
				var found bool
				hits.Lock()
				if ipr, found = hits.m[cdnIP.String()]; !found {
					ipr = &ipRate{}
					ipr.New(cdnIP)
					hits.m[cdnIP.String()] = ipr
				}
				hits.Unlock()
				overLimit := ipr.Hit()
				if overLimit {
					if whitelist.contains(cdnIP) {
						//fmt.Println("but is whitelisted so we don't care.")
					} else {
						if err := ipr.Limit(); err != nil {
							fmt.Printf("Error limiting IP: %s", err)
						}
					}
				}
			}
		}(channel)

		server.Wait()

		return nil
	}
	app.Run(os.Args)
}
<commit_msg>Track if an ipr was limited.<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/alienth/fastlyctl/util"
	"github.com/alienth/go-fastly"
	"github.com/juju/ratelimit"
	"github.com/urfave/cli"
	"gopkg.in/mcuadros/go-syslog.v2"
)

// getIPs takes in an haproxy log line and returns the client IP and the CDN
// connecting IP which is in the captured request header. For us, the CDN
// connecting IP happens to be the 8th captured request header.
func getIPs(logLine string) (*net.IP, *net.IP) {
	if logLine == "" {
		return nil, nil
	}
	// This string parsing stuff was lifted from TPS
	var a, b int
	if a = strings.Index(logLine, "]:") + 3; a == -1 {
		return nil, nil
	}
	if b = strings.Index(logLine[a:], ":"); b == -1 {
		return nil, nil
	}
	clientIPString := logLine[a : a+b]
	clientIP := net.ParseIP(clientIPString)
	if clientIP == nil {
		return nil, nil
	}

	logLine = logLine[a+b:]
	// The first curly-braced block contains our request headers
	if a = strings.Index(logLine, "{"); a == -1 {
		return &clientIP, nil
	}
	if b = strings.Index(logLine[a:], "}"); b == -1 {
		return &clientIP, nil
	}
	bracketedHeaders := logLine[a : a+b]
	headers := strings.Split(bracketedHeaders, "|")
	if len(headers) < 7 {
		return &clientIP, nil
	}
	ipString := headers[7]
	cdnIP := net.ParseIP(ipString)
	if cdnIP == nil {
		return &clientIP, nil
	}
	return &clientIP, &cdnIP
}

func getLogTimestamp(logLine string) int64 {
	timestampStr := strings.Split(logLine, " ")[6]
	ts, _ := time.Parse("[02/Jan/2006:15:04:05.999]", timestampStr)

	return (ts.Unix())
}

type ipList struct {
	ips  []net.IP
	nets []net.IPNet
}

func readIPList(filename string) (*ipList, error) {
	var list ipList
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Index(strings.TrimSpace(line), "#") == 0 {
			continue
		}
		if strings.TrimSpace(line) == "" {
			continue
		}
		if strings.Contains(line, "/") {
			_, n, err := net.ParseCIDR(line)
			if err != nil {
				return nil, fmt.Errorf("Unable to parse CIDR.\nLine:\n%s\nError:\n%s\n", line, err)
			}
			if n != nil {
				list.nets = append(list.nets, *n)
			}
		} else {
			ip := net.ParseIP(line)
			if ip != nil {
				list.ips = append(list.ips, ip)
			} else {
				return nil, fmt.Errorf("Unable to parse IP address in list: %s\n", line)
			}
		}

	}
	return &list, nil
}

func (l *ipList) contains(checkIP *net.IP) bool {
	if l == nil {
		return false
	}
	for _, ip := range l.ips {
		if ip.Equal(*checkIP) {
			return true
		}
	}
	for _, net := range l.nets {
		if net.Contains(*checkIP) {
			return true
		}
	}

	return false
}

type ipRate struct {
	ip      *net.IP
	bucket  *ratelimit.Bucket
	entries []*util.ACLEntry
	limited bool

	FirstHit    int64 `json:"first_hit,omitempty"`
	LastHit     int64 `json:"last_hit,omitempty"`
	LastLimit   int64 `json:"last_limit,omitempty"`
	Hits        int   `json:"hits,omitempty"`
	Strikes     int   `json:"strikes,omitempty"`
	Expire      int64 `json:"-"`
	LimitExpire int64 `json:"limit_expire,omitempty"`
}

func (ipr *ipRate) New(ip *net.IP) {
	ipr.ip = ip
	ipr.bucket = ratelimit.NewBucket(time.Duration(3)*time.Minute, 180)
	ipr.Expire = time.Now().Add(time.Duration(1) * time.Hour).Unix()
}

// Records a hit and returns true if it is over limit.
func (ipr *ipRate) Hit() bool {
	// Here we pretend that the IP performed no hits before the limit was removed,
	// since in testing the IP is not actually limited.
	if ipr.limited {
		return false
	}

	_, isSoonerThanMaxWait := ipr.bucket.TakeMaxDuration(1, 0)
	if ipr.FirstHit == 0 {
		ipr.FirstHit = time.Now().Unix()
	}
	ipr.LastHit = time.Now().Unix()
	ipr.Hits++
	ipr.Expire = time.Now().Add(time.Duration(1) * time.Hour).Unix()
	return !isSoonerThanMaxWait
}

// Limit adds an IP to a fastly edge ACL
func (ipr *ipRate) Limit() error {

	service, err := util.GetServiceByNameOrID(client, "teststackoverflow.com")
	if err != nil {
		return err
	}

	acl, err := util.NewACL(client, service.ID, "ratelimit")
	if err != nil {
		return err
	}

	// TODO this is heavy.. cache? do we even need?(see below)
	entries, err := acl.ListEntries()
	if err != nil {
		return err
	}

	var entry *util.ACLEntry

	for _, e := range entries {
		if ipr.ip.String() == e.IP {
			fmt.Printf("IP %s is already limited.\n", e.IP)
			// Not all ACL entries were necessarily created by this script, so silently ignore unmarshal errors.
			if err := json.Unmarshal([]byte(e.Comment), ipr); err != nil {
				break
			}
			entry = &util.ACLEntry{Client: client, ID: e.ID, ACLID: e.ACLID, ServiceID: service.ID}
			break
		}
	}

	ipr.LastLimit = time.Now().Unix()
	ipr.Strikes++
	ipr.LimitExpire = time.Now().Add(time.Duration(5) * time.Minute).Unix()
	ipr.Expire = time.Now().Add(time.Duration(24) * time.Hour).Unix()
	comment, err := json.Marshal(ipr)
	if err != nil {
		return err
	}
	if entry == nil {
		entry, err = util.NewACLEntry(client, service.ID, "ratelimit", ipr.ip.String(), 0, string(comment), false)
		if err != nil {
			return err
		}
		fmt.Printf("Limiting IP %s\n", ipr.ip.String())
		entry.Add()
	} else {
		entry.Comment = string(comment)
		//		entry.Update()
	}

	ipr.limited = true
	ipr.entries = append(ipr.entries, entry)
	return nil
}

// Removes an IP from ratelimits
func (ipr *ipRate) RemoveLimit() error {
	if len(ipr.entries) > 0 {
		fmt.Printf("Unlimiting IP %s\n", ipr.ip.String())
		for i, entry := range ipr.entries {
			if err := entry.Remove(); err != nil {
				fmt.Println(err)
			}
			ipr.entries[i] = nil
		}
		ipr.entries = ipr.entries[:0]
		ipr.limited = false
	}
	return nil
}

type hitMap struct {
	sync.RWMutex
	m map[string]*ipRate
}

func (hits *hitMap) expireRecords() {
	for {
		hits.Lock()
		for ip, ipr := range hits.m {
			if ipr.Expire < time.Now().Unix() {
				ipr.RemoveLimit()
				delete(hits.m, ip)
			}
		}
		hits.Unlock()
		time.Sleep(time.Duration(60) * time.Second)
	}
}

func (hits *hitMap) expireLimits() {
	for {
		hits.Lock()
		for _, ipr := range hits.m {
			if ipr.LimitExpire < time.Now().Unix() {
				ipr.RemoveLimit()
			}
		}
		hits.Unlock()
		time.Sleep(time.Duration(15) * time.Second)
	}
}

// Fetches down remote ACLs and populates local hitMap with previously stored data.
func (hits *hitMap) importIPRates() error {
	service, err := util.GetServiceByNameOrID(client, "teststackoverflow.com")
	if err != nil {
		return err
	}

	acl, err := util.NewACL(client, service.ID, "ratelimit")
	if err != nil {
		return err
	}

	entries, err := acl.ListEntries()
	if err != nil {
		return err
	}

	// TODO: work with IPs existing in multiple services.
	for _, e := range entries {
		var ipr *ipRate
		var found bool
		hits.Lock()
		if ipr, found = hits.m[e.IP]; !found {
			ipr = &ipRate{}
			ip := net.ParseIP(e.IP)
			if ip == nil {
				return fmt.Errorf("Unable to parse IP %s in ACL.")
			}
			ipr.New(&ip)
			hits.m[ip.String()] = ipr
		}
		hits.Unlock()
		if err := json.Unmarshal([]byte(e.Comment), ipr); err != nil {
			break
		}
		entry := &util.ACLEntry{Client: client, ID: e.ID, ACLID: e.ACLID, ServiceID: service.ID}
		ipr.limited = true
		ipr.entries = append(ipr.entries, entry)
	}

	return nil
}

var client *fastly.Client

func main() {
	app := cli.NewApp()
	app.Name = "fastly-ratelimit"

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "whitelist, w",
			Usage: "Read `FILE` containing addresses which should whitelisted from any bans.",
		},
		cli.StringFlag{
			Name:  "listen, l",
			Usage: "Specify listen `ADDRESS:PORT`.",
			Value: "0.0.0.0:514",
		},
		cli.StringFlag{
			Name:   "fastly-key, K",
			Usage:  "Fastly API Key. Can be read from 'fastly_key' file in CWD.",
			EnvVar: "FASTLY_KEY",
			Value:  util.GetFastlyKey(),
		},
	}
	app.Action = func(c *cli.Context) error {
		client, _ = fastly.NewClient(c.GlobalString("fastly-key"))
		channel := make(syslog.LogPartsChannel)
		handler := syslog.NewChannelHandler(channel)

		server := syslog.NewServer()
		server.SetFormat(syslog.RFC3164)
		server.SetHandler(handler)
		if err := server.ListenUDP(c.GlobalString("listen")); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to listen: %s\n", err), -1)
		}
		if err := server.Boot(); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to start server: %s\n", err), -1)
		}

		var whitelist *ipList
		var err error
		whitelistFile := c.GlobalString("whitelist")
		if whitelistFile != "" {
			if whitelist, err = readIPList(whitelistFile); err != nil {
				return cli.NewExitError(fmt.Sprintf("Error reading whitelist file: %s", err), -1)
			}
		}

		var hits = hitMap{m: make(map[string]*ipRate)}
		if err := hits.importIPRates(); err != nil {
			return cli.NewExitError(fmt.Sprintf("Error importing existing IP rates: %s", err), -1)
		}
		go hits.expireRecords()
		go hits.expireLimits()
		go func(channel syslog.LogPartsChannel) {
			for logParts := range channel {
				var line string
				var ok bool
				if line, ok = logParts["content"].(string); !ok || line == "" {
					continue
				}
				clientIP, cdnIP := getIPs(line)
				if cdnIP == nil || clientIP == nil {
					continue
				}
				var ipr *ipRate
				var found bool
				hits.Lock()
				if ipr, found = hits.m[cdnIP.String()]; !found {
					ipr = &ipRate{}
					ipr.New(cdnIP)
					hits.m[cdnIP.String()] = ipr
				}
				hits.Unlock()
				overLimit := ipr.Hit()
				if overLimit {
					if whitelist.contains(cdnIP) {
						//fmt.Println("but is whitelisted so we don't care.")
					} else {
						if err := ipr.Limit(); err != nil {
							fmt.Printf("Error limiting IP: %s", err)
						}
					}
				}
			}
		}(channel)

		server.Wait()

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


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strings"
)

const commentChar string = "#"

type hostsFileLine struct {
	ip     string
	hosts  []string
	lineNo int
}

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

func isComment(line string) bool {
	trimLine := strings.TrimSpace(line)
	isComment := strings.HasPrefix(trimLine, commentChar)
	return isComment
}

func itemInSlice(item string, list []string) bool {
	for _, i := range list {
		if i == item {
			return true
		}
	}

	return false
}

func getHostLines() []hostsFileLine {
	var output []hostsFileLine

	hostsFile, err := os.Open("/etc/hosts")
	check(err)
	defer hostsFile.Close()

	lineNo := 0

	scanner := bufio.NewScanner(hostsFile)
	for scanner.Scan() {
		line := scanner.Text()
		if !isComment(line) {
			fields := strings.Fields(line)
			if len(fields) > 0 {
				line := hostsFileLine{ip: fields[0], hosts: fields[1:], lineNo: lineNo}
				output = append(output, line)
			}
		}
		lineNo += 1
	}

	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}

	return output
}

func list() {
	lines := getHostLines()
	for i := range lines {
		line := lines[i]
		fmt.Printf("%s %s\n", line.ip, line.hosts)
	}
}

func lineInHosts(ip string, host string) bool {
	lines := getHostLines()
	for i := range lines {
		if ip == lines[i].ip {
			if itemInSlice(host, lines[i].hosts) {
				return true
			}
		}
	}

	return false
}

func addLine(ip string, host string) {
	fmt.Printf("About to add %s and %s to the hosts file.\n", ip, host)
}

func main() {
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "list":
			list()
			return
		case "check":
			if len(os.Args) < 3 {
				fmt.Println("usage: goodhost check 127.0.0.1 salt")
				os.Exit(1)
			}

			ip := os.Args[2]
			host := os.Args[3]
			if !lineInHosts(ip, host) {
				fmt.Fprintf(os.Stderr, "%s %s is not in the hosts file\n", ip, host)
				os.Exit(1)
			}
			return
		case "add":
			if len(os.Args) < 3 {
				fmt.Println("usage: goodhost add 127.0.0.1 myhost")
				os.Exit(1)
			}
			ip := os.Args[2]
			host := os.Args[3]

			addLine(ip, host)
			return
		}
	}

	fmt.Println("Help should go here.")
}
<commit_msg>Added rough first draft of add command.<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"os/user"
	"strings"
)

const commentChar string = "#"

type hostsFileLine struct {
	ip     string
	hosts  []string
	lineNo int
}

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

func isComment(line string) bool {
	trimLine := strings.TrimSpace(line)
	isComment := strings.HasPrefix(trimLine, commentChar)
	return isComment
}

func itemInSlice(item string, list []string) bool {
	for _, i := range list {
		if i == item {
			return true
		}
	}

	return false
}

func getHostLines() []hostsFileLine {
	var output []hostsFileLine

	hostsFile, err := os.Open("/etc/hosts")
	check(err)
	defer hostsFile.Close()

	lineNo := 0

	scanner := bufio.NewScanner(hostsFile)
	for scanner.Scan() {
		line := scanner.Text()
		if !isComment(line) {
			fields := strings.Fields(line)
			if len(fields) > 0 {
				line := hostsFileLine{ip: fields[0], hosts: fields[1:], lineNo: lineNo}
				output = append(output, line)
			}
		}
		lineNo += 1
	}

	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}

	return output
}

func list() {
	lines := getHostLines()
	for i := range lines {
		line := lines[i]
		fmt.Printf("%s %s\n", line.ip, line.hosts)
	}
}

func lineInHosts(ip string, host string) bool {
	lines := getHostLines()
	for i := range lines {
		if ip == lines[i].ip {
			if itemInSlice(host, lines[i].hosts) {
				return true
			}
		}
	}

	return false
}

func addLine(ip string, host string) {
	var lines []string

	hostsFile, err := os.OpenFile("/etc/hosts", os.O_APPEND|os.O_RDWR, 0770)
	check(err)
	defer hostsFile.Close()

	lineNo := 0
	lastLineNo := -1

	scanner := bufio.NewScanner(hostsFile)
	for scanner.Scan() {
		line := scanner.Text()
		lines = append(lines, line)
		if !isComment(line) {
			fields := strings.Fields(line)
			if len(fields) > 0 {
				if ip == fields[0] {
					lastLineNo = lineNo
				}
			}
		}
		lineNo += 1
	}
	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}

	if lastLineNo == -1 {
		// Ip line is already in file, so we just append our new line.
		endLine := fmt.Sprintf("%s %s", ip, host)
		lines = append(lines, endLine)
	} else {
		// Otherwise, we replace the line in the correct position
		lines[lastLineNo] = fmt.Sprintf("%s %s", lines[lastLineNo], host)
	}

	newHostsFile, err := os.Create("/etc/hosts")
	check(err)

	w := bufio.NewWriter(newHostsFile)
	defer w.Flush()

	for _, line := range lines {
		fmt.Fprintln(w, line)
	}
}

func main() {
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "list":
			list()
			return
		case "check":
			if len(os.Args) < 3 {
				fmt.Println("usage: goodhost check 127.0.0.1 salt")
				os.Exit(1)
			}

			ip := os.Args[2]
			host := os.Args[3]
			if !lineInHosts(ip, host) {
				fmt.Fprintf(os.Stderr, "%s %s is not in the hosts file\n", ip, host)
				os.Exit(1)
			}
			return
		case "add":
			if len(os.Args) < 3 {
				fmt.Println("usage: goodhost add 127.0.0.1 myhost")
				os.Exit(1)
			}
			user, err := user.Current()
			check(err)

			ip := os.Args[2]
			host := os.Args[3]
			if lineInHosts(ip, host) {
				fmt.Fprintf(os.Stderr, "Line already in host file. Nothing to do.\n")
				os.Exit(2)
			}

			if user == nil || user.Uid != "0" {
				fmt.Fprintf(os.Stderr, "Need to be root user. Try running with sudo.\n")
				os.Exit(1)
			}

			addLine(ip, host)
			return
		}
	}

	fmt.Println("Help should go here.")
}
<|endoftext|>