From 36fba6ff541ea2fec692489c9558944acff0a112 Mon Sep 17 00:00:00 2001
From: Martin Polden <martin.polden@gmail.com>
Date: Sat, 13 Aug 2016 20:32:19 +0200
Subject: [PATCH] Use logrus for logging

---
 api/api.go | 18 +++++++++---------
 main.go    | 13 +++++++++++--
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/api/api.go b/api/api.go
index 9f44c0f..dacc464 100644
--- a/api/api.go
+++ b/api/api.go
@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"html/template"
 	"io"
-	"log"
+
+	"github.com/Sirupsen/logrus"
+
 	"math/big"
 	"net"
 	"net/http"
@@ -27,6 +29,7 @@ type API struct {
 	Template string
 	IPHeader string
 	oracle   Oracle
+	log      *logrus.Logger
 }
 
 type Response struct {
@@ -43,8 +46,8 @@ type PortResponse struct {
 	Reachable bool   `json:"reachable"`
 }
 
-func New(oracle Oracle) *API {
-	return &API{oracle: oracle}
+func New(oracle Oracle, logger *logrus.Logger) *API {
+	return &API{oracle: oracle, log: logger}
 }
 
 func ipToDecimal(ip net.IP) *big.Int {
@@ -81,15 +84,15 @@ func (a *API) newResponse(r *http.Request) (Response, error) {
 	ipDecimal := ipToDecimal(ip)
 	country, err := a.oracle.LookupCountry(ip)
 	if err != nil {
-		log.Print(err)
+		a.log.Debug(err)
 	}
 	city, err := a.oracle.LookupCity(ip)
 	if err != nil {
-		log.Print(err)
+		a.log.Debug(err)
 	}
 	hostnames, err := a.oracle.LookupAddr(ip)
 	if err != nil {
-		log.Print(err)
+		a.log.Debug(err)
 	}
 	return Response{
 		IP:        ip,
@@ -211,9 +214,6 @@ type appHandler func(http.ResponseWriter, *http.Request) *appError
 
 func (fn appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if e := fn(w, r); e != nil { // e is *appError
-		if e.Error != nil {
-			log.Print(e.Error)
-		}
 		// When Content-Type for error is JSON, we need to marshal the response into JSON
 		if e.IsJSON() {
 			var data = struct {
diff --git a/main.go b/main.go
index ed8e762..42415da 100644
--- a/main.go
+++ b/main.go
@@ -3,9 +3,10 @@ package main
 import (
 	flags "github.com/jessevdk/go-flags"
 
-	"log"
 	"os"
 
+	"github.com/Sirupsen/logrus"
+
 	"github.com/martinp/ipd/api"
 )
 
@@ -18,12 +19,20 @@ func main() {
 		PortLookup    bool   `short:"p" long:"port-lookup" description:"Enable port lookup"`
 		Template      string `short:"t" long:"template" description:"Path to template" default:"index.html"`
 		IPHeader      string `short:"H" long:"trusted-header" description:"Header to trust for remote IP, if present (e.g. X-Real-IP)"`
+		LogLevel      string `short:"L" long:"log-level" description:"Log level to use" default:"info" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic"`
 	}
 	_, err := flags.ParseArgs(&opts, os.Args)
 	if err != nil {
 		os.Exit(1)
 	}
 
+	log := logrus.New()
+	level, err := logrus.ParseLevel(opts.LogLevel)
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Level = level
+
 	oracle := api.NewOracle()
 	if opts.ReverseLookup {
 		log.Println("Enabling reverse lookup")
@@ -49,7 +58,7 @@ func main() {
 		log.Printf("Trusting header %s to contain correct remote IP", opts.IPHeader)
 	}
 
-	api := api.New(oracle)
+	api := api.New(oracle, log)
 	api.Template = opts.Template
 	api.IPHeader = opts.IPHeader