Host checker

hostchecker.go

/* TCP Host checker 
 * Eric Belanger // github.com/bilange // Twitter: @bilange
 *
 * I needed a quick diagnostic tool to troubleshoot a specific issue with my
 * Ubuntu installation at home: apparently after a long, sustained TCP
 * connection (thing SSH or VPN), somehow the system throws its hands in the
 * air and refuses to talk to the same HOST, until I reboot my PC. 
 *
 * I basically needed to know when a remote host would stop responding to my
 * requests. Creating a Go program could be seen as overkill, but I wanted to
 * push on my Go knowledge, and thought it was a very good opportunity.
 *
 * I have commented my code for peer review and for learning purposes.
 *
 * This MAY be NOT the only way to do this task (especially when you think
 * about the "flag" Go package), so feel free to fork, base upon and
 * experiment. Hack away! :D
 *
 * Technical note: I originally wanted to include an optional ICMP Echo (ping)
 * method of checking, however it seems sending an ICMP echo request somehow
 * needs root (google "golang ping" for the juicy details). As I intent to
 * deploy this on non-root accounts for scripting purposes, i dropped this
 * 'native' programatically generated ping feature. I could however launch a
 * subprocess calling the OS' "ping" command and check for the result, however
 * this is highly operating system's dependant (and very variable). I leave
 * that as an exercice for the reader.
 */ 
package main

import (
  "fmt"
  "strconv"
  "time"
  "net"
  "os"
  "os/signal"
)

//Note that the variables values here reflects the 'default' behaviour, until
//the caller says otherwise when invoking the command parameters

var Host string = "apple.com" //Where we should connect to (host)
var Port int = 80             //Where we should connect to (port)
var Times int = -1            //How many times we're checking (<= 0 == infinite)
var TimesDone int = 0         //Counter keeping track how many times we checked so far
var Delay int = 300           //Interval, in seconds, between each checks.

var ExitCodeSuccess int = 0   //-r sets the exit code on host response (success), for scripting
var ExitCodeFailure int = 1   //-n sets the exit code on host non-response (failure), for scripting

var Verbose bool = false      //Should we print out stuff on screen?
var AlmostQuiet bool = false  //Only print "host is up" / "host is down".
var HostIsUp bool = false     //Was the host up at the last checkup?

//This function is called for every integer based commandline parameter we need
//to parse.  If parsing of variable v fails, integer d will be used.
func parseInt(v string, d int) int {
  rtn64, err := strconv.ParseInt(v, 10, 32)
  if err != nil { return d }
  return int(rtn64)
}

func usage() {
  fmt.Printf("Usage: %s [options] <target IP or hostname>\nwhere <options> is any combination of these:\n\n%s\n", os.Args[0], 
`-p, --port            Specifies a TCP port to check (Ex.: -p 22)
                      Defaults to 80 (HTTP)

-r, --response        A succesful check will return this exit code (Ex.: -r 0)
-n, --noresponse      An unsuccesful check will return this exit code (Ex.: -r 1)

-t, --times           Number of checks done. Omit this parameter will 
                      cause one check, for scripting purposes. Set
                      this to 0 or less for an infinite check. Any multiple
                      check amount will enable verbosity on screen.
-d, --delay           Delays every check by that many seconds (Ex.: -d 60)

-v, --verbose         Force verbosity when doing one single check
-a, --almostquiet     Only print "host is down" / "host is up"

By default, this program will check for a TCP connection on apple.com:80, 
every five minutes, printing results on screen. 
`)
  os.Exit(0)
}

//Shortcut function for verbose printing, displaying the time in the 
//format "[Sun Oct 28 15:43:33 2012]"
func timestamp() string {
  return "["+time.Now().Format(time.ANSIC)+"] "
}

//Returns true if we need to print out on screen
//printInAlmostQuietMode is basically an override that permits printing, whether
//we're in verbose mode or not.
func stdoutRequired() bool {
  switch {
    //case Times > 1 && AlmostQuiet == true: return true
    case Verbose == true && Times == 1: return true
    case Verbose == false && Times == 1: return false

    case Verbose == true && Times != 1: return true
    case Verbose == false && Times != 1: return false
  }
  return false
}

func main() {
  //If the user hasnt specified ANYTHING, spit out the usage documenation and exit
  if len(os.Args) == 1 { usage() }

  //Parse the commandline parameters. I used my own parser instead of the
  //provided "flag" Go package because the parameters are optional here, and I
  //needed one single argument without being a "dashed" option. (The hostname)
  for i := 1; i<len(os.Args); i++ {
    var arg = os.Args[i]
    switch arg {
      case "-h","--help":
        usage()
      case "-p","--port":
        i++
        if i == len(os.Args)-1 { continue } //Reached end of the slice
        Port = parseInt(os.Args[i], 80)
        continue
      case "-v","--verbose":
        Verbose = true
        fmt.Println(timestamp(), "Verbose mode on.")
      case "-a","--almostquiet":
        AlmostQuiet = true
      case "-t","--times":
        i++
        if i == len(os.Args)-1 { continue } //Reached end of the slice
        Times = parseInt(os.Args[i], 80)
        continue
      case "-d","--delay":
        i++
        if i == len(os.Args)-1 { continue } //Reached end of the slice
        Delay = parseInt(os.Args[i], 300)
        continue
      case "-r","--response":
        i++
        if i == len(os.Args)-1 { continue } //Reached end of the slice
        ExitCodeSuccess = parseInt(os.Args[i], 0)
        continue
      case "-n","--noresponse":
        i++
        if i >= len(os.Args) { continue } //Reached end of the slice
        ExitCodeFailure = parseInt(os.Args[i], 1)
        continue
      default: 
        Host = arg
    }
  }
  
  //Safe checking of incompatible settings
  if Times != 1 && Verbose == false {
    fmt.Println(timestamp(), "Non-single check specified without verbosity; a brief message will appear for every check.")
    AlmostQuiet = true
  }
  if AlmostQuiet == true && Verbose == true {
    fmt.Println(timestamp(), "Ambiguous parameters: AlmostQuiet AND verbose specified: Verbose wins.")
    AlmostQuiet = false
  }
  if (Delay < 30) { 
    if AlmostQuiet == false {
      fmt.Println(timestamp(), "** WARNING ** Causing too much traffic is usually frown upon by sysadmins. Use with care!")
    }
  }
  if stdoutRequired() && AlmostQuiet == false { 
    fmt.Printf("%s Settings used: \n\nHost: %s, Port: %d\nSuccess Exit Code: %d, Failure Exit Code: %d\nDelay: %ds, Times Checked: %d, Verbose: %t, Almost quiet: %t\n\n",
    timestamp(), Host, Port, ExitCodeSuccess, ExitCodeFailure, Delay, Times, Verbose, AlmostQuiet) 
  }

  //This part of code is heavily inspired from:
  //http://stackoverflow.com/questions/11268943/golang-is-it-possible-to-capture-a-ctrlc-sigterm-signal-and-run-a-cleanup-fu
  c := make (chan os.Signal, 1)
  signal.Notify(c, os.Interrupt)
  go func(){
    for sig := range c {
      if stdoutRequired() { fmt.Printf("^C hit, ") }
      switch HostIsUp {
        case true: 
          if AlmostQuiet == true || stdoutRequired() { fmt.Printf("host is UP.\n" ) }
          os.Exit(ExitCodeSuccess)
        case false: 
          if AlmostQuiet == true || stdoutRequired() { fmt.Printf("host is DOWN.\n") }
          os.Exit(ExitCodeFailure)
      }
      if stdoutRequired() { fmt.Printf("Exiting. (%v)\n", sig) }
    }
  }()

  for { //Loops until we have reached the amount of times the user wanted to check up the host.
    //Note that exiting is handled just above this for loop.

    if stdoutRequired() == true { fmt.Printf("%sAttempting to connect to %s:%d... ",timestamp(), Host, Port) }
    if Times > 0 { TimesDone++ } //Only keep track of the times we check when we need it :)

    conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", Host, Port))
    if err != nil { //An error indicates that net.Dial couldn't reach the host.
      HostIsUp = false
    } else {
      if stdoutRequired() { fmt.Printf("Connected. ") }
      conn.Close()
      HostIsUp = true
    }

    if HostIsUp == true {
      if AlmostQuiet == true || stdoutRequired() { fmt.Printf("Host is UP!\n") }
      if Times > 0 && TimesDone >= Times { //Finite check
          os.Exit(ExitCodeSuccess)
      } 
    }
    if HostIsUp == false {
      if AlmostQuiet == true || stdoutRequired() { fmt.Printf("Host is DOWN!\n") }
      if Times > 0 && TimesDone >= Times { //Finite check
          os.Exit(ExitCodeFailure)
      } 
    }

    time.Sleep(time.Duration(Delay) * time.Second) 
  }
}