#!/usr/bin/ruby1.8
#
# This daemon will compute current network usage statistics:
#   - today's usage
#   - today's limit
#
# It will publish these to the file /var/run/gatekeeper/stats
# which can be read by another program to display the stats.
# When today's usage exceeds today's limit, this will disconnect
# the local squid cache from the remote squid cache.

require 'netusage'

require "yaml"
require "pathname"
require "fileutils"

CONFIGURATION_FILE = '/etc/gatekeeper.conf'
begin ; CONFIGURATION = YAML::load( File.open( CONFIGURATION_FILE ) ) ; rescue ; raise "Unable to load #{CONFIGURATION_FILE}" ; end

def connect
  # Establish the connection between local and remote squids:
  return system( CONFIGURATION['ip_connect'] )
end

def disconnect
  # Cut the connection between local and remote squids:
  return system( CONFIGURATION['ip_disconnect'] )
end

def publish (usage, limit)
  # Make the current usage available for display.
  begin
#    FileUtils::mkpath("/var/run/gatekeeper")
    File.open( CONFIGURATION['statsFile'], "w") do |pubFile|
      YAML.dump({"Usage" => usage, "Limit" => limit}, pubFile)
    end
  rescue
    # This must not throw an exception.
  end

end

def getDate
  # Return the date as the number of days since 1970.
  # If we decide to have the day begin at 6 a.m. instead of midnight,
  # we can do that here by subtracting 6*60*60 from Time.now.to_f.
  return (Time.now.to_f / (24 * 60 * 60)).floor
end

def saveState (todaysDate, todaysUsage, todaysLimit)
  # Save the gatekeeper's state so that it works across
  # restarts and reboots.
  state = {}
  state["todaysDate"] = todaysDate
  state["todaysUsage"] = todaysUsage
  state["todaysLimit"] = todaysLimit
  begin
#    FileUtils::mkpath(Pathname.new(CONFIGURATION['stateFilePath']).dirname.to_s)
    File.open(CONFIGURATION['stateFilePath'], "w") do |stateFile|
      YAML.dump(state, stateFile)
    end
  rescue
  end
end

def initState
  # Read the state information from the state file, creating the file if necessary.
  begin
    File.open(CONFIGURATION['stateFilePath']) do |stateFile|
      return YAML.load(stateFile)
    end
  rescue
    # If the file doesn't exist, create it using an initial state:
    todaysDate = getDate
    todaysUsage = 0
    todaysLimit = CONFIGURATION['dailyLimit']
    saveState(todaysDate, todaysUsage, todaysLimit)
    begin
      File.open(CONFIGURATION['stateFilePath']) do |stateFile|
        return YAML.load(stateFile)
      end
    rescue
      return {"todaysDate" => todaysDate, "todaysUsage" => todaysUsage, "todaysLimit" => todaysLimit}
    end
  end
end

def gateKeeper (iface)
  # Periodically compute and publish the day's usage and current limit.
  # If it is over the limit
  #   call /etc/iptables/disconnect
  # otherwise
  #   call /etc/iptables/connect
  # Recompute the usage and limit every morning.

  # On startup, initialize state information:
  state = initState
  todaysDate = state["todaysDate"]
  todaysUsage = state["todaysUsage"]
  todaysLimit = state["todaysLimit"]

  # Use a CalcDelta instance to compute network usage.
  cd = NetUsage::CalcDelta.new(iface)

  while true
    # Update variables in the morning:
    date = getDate
    if date > todaysDate
      todaysLimit = CONFIGURATION['dailyLimit'] * (date - todaysDate) + todaysLimit - todaysUsage
      if todaysLimit < 0
        todaysLimit = 0
      else if CONFIGURATION['maxLimit'] < todaysLimit
             todaysLimit = CONFIGURATION['maxLimit']
           end
      end
      todaysDate = date
      todaysUsage = 0
    end

    todaysUsage += cd.usage

    publish todaysUsage, todaysLimit
    saveState(todaysDate, todaysUsage, todaysLimit)

    if todaysUsage <= todaysLimit
      connect
    else
      disconnect
    end

    sleep 1
  end
end

gateKeeper(CONFIGURATION['monitorInterface'])