get_shit_done.rb

#!/usr/bin/env ruby

require "fileutils"
require "set"

class GetShitDone
  attr_accessor :hosts_file
  def initialize
    @hosts_file = "/etc/hosts"
    @banned_sites_file = "#{ENV['HOME']}/.banned_sites"
    @start_hosts_separator = "--- start-get-shit-done.rb ---"
    @end_hosts_separator = "--- end-get-shit-done.rb ---"
  end

  def run
    loop do
      check_and_set_mode
      sleep(600)
    end
  end

  def log(message)
    STDOUT.puts(message)
  end

  def check_and_set_mode
    current_hour = Time.now.hour
    log("Checking time and current time is #{current_hour}")
    if current_hour > 10 && current_hour < 19
      set_work_mode
    else
      log("Setting play mode, at #{Time.now}")
      remove_banned_sites
    end
    system("dscacheutil -flushcache")
  end

  def set_work_mode
    log "Setting work mode now, #{Time.now}"
    read_banned_sites
    add_banned_sites
  end

  def current_mode
    return @work_mode if @work_mode
    work_mode_flag = false
    File.open(@hosts_file,"r") do |fl|
      lines = fl.readlines()
      work_mode_flag = lines.any? { |line| line =~ /#{@start_hosts_separator}/}
    end
    @work_mode = work_mode_flag ? :work : :play
  end

  def work_mode?
    @work_mode == :work
  end

  def play_mode?
    @work_mode == :play
  end

  def read_banned_sites
    fl = File.open(@banned_sites_file,"r")
    @sites = Set.new()
    fl.readlines.each do |line|
      line.strip!
      if line && !line.empty?
        @sites << "127.0.0.1 \t #{line}\n"
        @sites << "127.0.0.1 \t www.#{line}\n"
      end
    end
  ensure
    fl.close()
  end

  def add_banned_sites
    normal_host_lines = regular_host_lines
    File.open(@hosts_file, "w") do |fl|
      fl.write normal_host_lines.join()
      fl.puts
      fl.puts @start_hosts_separator
      fl.puts @sites.to_a.join()
      fl.puts @end_hosts_separator
    end
  end

  def regular_host_lines
    lines = Set.new()
    start_skipping = false

    File.open(@hosts_file, "r") do |fl|
      fl.readlines.each do |line|
        if line =~ /#{@start_hosts_separator}/
          start_skipping = true
          next
        end

        if line =~ /#{@end_hosts_separator}/
          start_skipping = false
          next
        end

        unless start_skipping
          lines << line
        end
      end
    end
    lines.to_a
  end

  def remove_banned_sites
    lines = regular_host_lines
    fl = File.open(@hosts_file,"w")
    fl.write(lines.join())
  ensure
    fl.close()
  end

end


if __FILE__ == $0
  if fork
    sleep(5)
    exit(0)
  else
    Process.setsid
    STDOUT.reopen("/var/log/get_shit_done.log")
    STDOUT.sync = true
    GetShitDone.new.run()
  end
end