frenzy.rb

#!/usr/bin/env ruby -rubygems
# frenzy - a script for adding/removing /etc/hosts site-blocking entries

require 'main'
require 'yaml'

SITES_TEMPLATE = <<EOF.strip
blocked_sites:
  - *.reddit.com
  - www.reddit.com
  - reddit.com
  - programming.reddit.com
  - news.ycombinator.com
  - twitter.com
  - www.twitter.com
  - tumblr.com
  - www.tumblr.com
  - facebook.com
  - www.facebook.com
  - *.facebook.com
  #- boingboing.net
EOF

FRENZY_START = '### BEGIN FRENZY BLOCK'
FRENZY_END = '### END FRENZY BLOCK'
FRENZY_MESSAGE = <<EOF.strip
# This section will be changed when turning frenzy on and off. When
# frenzy changes this block, it will first back up /etc/hosts to 
# /etc/hosts.bak_frenzy.
EOF

# TODO: implement commenting/uncommenting and add this to the comments
#
# To add custom lines that will be commented and uncommented by
# frenzy, add ##CUSTOM to the end of them. (It can appear after
# a comment.)

HELP_MESSAGE = <<EOF.strip
To enable turning frenzy mode on and off, add the following to /etc/hosts:

#{FRENZY_START}
#{FRENZY_MESSAGE}
#{FRENZY_END}

BLOCKED SITES
=============

The blocked site list is in ~/.frenzy/sites.yaml.
It gets added when you run "#{$0} setup".
You can run "#{$0} off" when the frenzy is off and it will
re-apply the current blocked sites list.

COMMANDS
========

#{$0}       - print the status
#{$0} setup - create ~/.frenzy/sites.yaml (blocked site list)
#{$0} on    - turn site-blocking on (run with sudo)
#{$0} off   - turn site-blocking off (run with sudo)
#{$0} help  - show this help message
EOF

class Frenzy
  def initialize()
  end

  def check_setup()
    unless File.exists?(File.expand_path('~/.frenzy'))
      puts "To set up frenzy, run '#{$0} setup'."
      exit
    end
  end

  def setup()
    FileUtils.mkdir File.expand_path('~/.frenzy')
    File.open(File.expand_path('~/.frenzy/sites.yml'), 'w') do |f|
      f.puts SITES_TEMPLATE
    end
    puts HELP_MESSAGE
    puts
    puts "To show this help message again, run '#{$0} help'."
  end

  def print_status()
    if File.read('/etc/hosts').include?('frenzy is ON')
      puts 'The frenzy is ON!'
    elsif File.read('/etc/hosts').include?('frenzy is OFF')
      puts 'The frenzy is OFF.'
    else
      puts 'The frenzy status was not found in /etc/hosts.'
    end
  end

  def print_help()
    puts HELP_MESSAGE
  end

  def check_hosts_writable()
    unless File.writable?('/etc/hosts')
      puts "Cannot write to /etc/hosts. Exiting."
    end
  end

  def backup_hosts()
    FileUtils.cp('/etc/hosts', '/etc/hosts.bak_frenzy')
  end

  def render_entries(f)
    sites = YAML::load_file(File.expand_path('~/.frenzy/sites.yml'))
    sites['blocked_sites'].each do |site|
      f.puts "127.0.0.1\t\t#{site}"
    end
  end

  def turn(state)
    check_hosts_writable
    backup_hosts
    copying = true
    inserted = false
    File.open('/etc/hosts.new_frenzy', 'w') do |w|
      File.open('/etc/hosts') do |r|
        r.each_line do |line|
          if line.strip == FRENZY_START
            copying = false
            w.puts FRENZY_START
            w.puts FRENZY_MESSAGE
            w.puts '# The frenzy is ' + state.to_s.upcase
            render_entries w if state == :off
            inserted = true
          elsif line.strip == FRENZY_END
            w.puts FRENZY_END
            copying = true
          else
            w.puts line if copying
          end
        end
      end
    end
    if copying and inserted
      FileUtils.cp '/etc/hosts.new_frenzy', '/etc/hosts'
    else
      puts 'The required comments to guide the script were not found in /etc/hosts.'
      puts "Run #{$0} help for instructions on inserting the required snippets."
    end
    FileUtils.rm '/etc/hosts.new_frenzy'

    system 'lookupd -flushcache'
  end

  def turn_on()
    turn :on
  end

  def turn_off()
    turn :off
  end
end

Main {
  mode 'setup' do
    def run()
      frenzy = Frenzy.new
      frenzy.setup
    end
  end

  mode 'help' do
    def run()
      frenzy = Frenzy.new
      frenzy.print_help
    end
  end

  mode 'on' do
    def run()
      frenzy = Frenzy.new
      frenzy.check_setup
      frenzy.turn_on
      frenzy.print_status
    end
  end

  mode 'off' do
    def run()
      frenzy = Frenzy.new
      frenzy.check_setup
      frenzy.turn_off
      frenzy.print_status
    end
  end

  def run()
    frenzy = Frenzy.new
    frenzy.check_setup
    frenzy.print_status
  end
}