Cacher is the code snippet organizer for pro developers

#!/usr/bin/env ruby -w

# brew-services(1) - Easily start and stop formulas via launchctl
# ===============================================================
#
# ## SYNOPSIS
#
# [<sudo>] `brew services` `list`<br>
# [<sudo>] `brew services` `restart` <formula><br>
# [<sudo>] `brew services` `start` <formula> [<plist>]<br>
# [<sudo>] `brew services` `stop` <formula><br>
# [<sudo>] `brew services` `cleanup`<br>
#
# ## DESCRIPTION
#
# Integrates homebrew formulas with MacOS X' `launchctl` manager. Services
# can either be added to `/Library/LaunchDaemons` or `~/Library/LaunchAgents`.
# Basically items added to `/Library/LaunchDaemons` are started at boot,
# those in `~/Library/LaunchAgents` at login.
#
# When started with `sudo` it operates on `/Library/LaunchDaemons`, else
# in the user space.
#
# Basically on `start` the plist file is generated and written to a `Tempfile`,
# then copied to the launch path (existing plists are overwritten).
#
# ## OPTIONS
#
# To access everything quickly, some aliases have been added:
#
#  * `rm`:
#    Shortcut for `cleanup`, because that's basically whats being done.
#
#  * `ls`:
#    Because `list` is too much to type :)
#
#  * `reload', 'r':
#    Alias for `restart`, which gracefully restarts selected service.
#
#  * `load`, `s`:
#    Alias for `start`, guess what it does...
#
#  * `unload`, `term`, `t`:
#    Alias for `stop`, stops and unloads selected service.
#
# ## SYNTAX
#
# Several existing formulas (like mysql, nginx) already write custom plist
# files to the formulas prefix. Most of these implement `#startup_plist`
# which then in turn returns a neat-o plist file as string.
#
# `brew services` operates on `#startup_plist` as well and requires
# supporting formulas to implement it. This method should either string
# containing the generated XML file, or return a `Pathname` instance which
# points to a plist template, or a hash like:
#
#    { :url => "https://gist.github.com/raw/534777/63c4698872aaef11fe6e6c0c5514f35fd1b1687b/nginx.plist.xml" }
#
# Some simple template parsing is performed, all variables like `{{name}}` are
# replaced by basically doing:
# `formula.send('name').to_s if formula.respond_to?('name')`, a bit like
# mustache. So any variable in the `Formula` is available as template
# variable, like `{{var}}`, `{{bin}}` usw.
#
# ## EXAMPLES
#
# Install and start service mysql at boot:
#
#     $ brew install mysql
#     $ sudo brew services start mysql
#
# Stop service mysql (when launched at boot):
#
#     $ sudo brew services stop mysql
#
# Start memcached at login:
#
#     $ brew install memcached
#     $ brew services start memcached
#
# List all running services for current user, and root:
#
#     $ brew services list
#     $ sudo brew services list
#
# ## BUGS
#
# `brew-services.rb` might not handle all edge cases, though it tries
# to fix problems by running `brew services cleanup`.
#
# ## COPYRIGHT
#
# Copyright (c) 2010 Lukas Westermann <lukas@at-point.ch>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
module ServicesCli
  class << self      
    # Binary name.
    def bin; "brew services" end

    # Path to launchctl binary.
    def launchctl; "/bin/launchctl" end

    # Wohoo, we are root dude!
    def root?; Process.uid == 0 end

    # Current user, i.e. owner of `HOMEBREW_CELLAR`.
    def user; @user ||= %x{/usr/bin/stat -f '%Su' #{HOMEBREW_CELLAR} 2>/dev/null}.chomp || %x{/usr/bin/whoami}.chomp end

    # Run at boot.
    def boot_path; Pathname.new("/Library/LaunchDaemons") end

    # Run at login.
    def user_path; Pathname.new(ENV['HOME'] + '/Library/LaunchAgents') end

    # If root returns `boot_path` else `user_path`.
    def path; root? ? boot_path : user_path end

    # Print warning message, which looks similar to homebrew's.
    def eww(msg); STDERR.puts("#{Tty.red}Warning#{Tty.reset}: #{msg}"); true end

    # Display error message and fail using `abort` if `fail == true`.
    def argh(msg, fail = true); "#{Tty.red}Error#{Tty.reset}: #{msg}".tap { |s| fail ? abort(s) : STDERR.puts(s) } end

    # Find all currently running services via launchctl list
    def running; %x{#{launchctl} list | grep com.github.homebrew}.chomp.split("\n").map { |svc| $1 if svc =~ /(com\.github\.homebrew\..+)\z/ }.compact end

    # Check if running as homebre and load required libraries et al.
    def homebrew!
      abort("Runtime error: homebrew is required, please start via `#{bin} ...`") unless defined?(HOMEBREW_LIBRARY_PATH)      
      %w{fileutils pathname tempfile formula}.each { |req| require(req) }
      self.send(:extend, ::FileUtils)
      ::Formula.send(:include, Service::PlistSupport)
    end

    # Access current service
    def service; @service ||= Service.new(Formula.factory(Formula.resolve_alias(@formula))) if @formula end

    # Print usage and `exit(...)` with supplied exit code, if code
    # is set to `false`, then exit is ignored.
    def usage(code = 0)
      puts "usage: [sudo] #{bin} [--help] <command> [<formula>]"
      puts     
      puts "Small wrapper around `launchctl` for supported formulas, commands available:"
      puts "   cleanup Get rid of stale services and unused plists"
      puts "   list    List all services managed by `#{bin}`"
      puts "   restart Gracefully restart selected service"
      puts "   start   Start selected service"
      puts "   stop    Stop selected service"
      puts
      puts "Options, sudo and paths:"
      puts
      puts "  sudo   When run as root, operates on #{boot_path} (run at boot!)"
      puts "  Run at boot:  #{boot_path}"
      puts "  Run at login: #{user_path}"
      puts
      exit(code) unless exit == false
      true
    end
    
    # Run and start the command loop.
    def run!    
      homebrew!
      usage if ARGV.empty? || ARGV.include?('help') || ARGV.include?('--help') || ARGV.include?('-h')

      # parse arguments
      @args = ARGV.reject { |arg| arg[0] == 45 }.map { |arg| arg.include?("/") ? arg : arg.downcase } # 45.chr == '-'
      @cmd = @args.shift
      @formula = @args.shift
      
      # dispatch commands and aliases
      case @cmd
        when 'cleanup', 'clean', 'cl', 'rm' then cleanup
        when 'list', 'ls' then list
        when 'restart', 'relaunch', 'reload', 'r' then check and restart
        when 'start', 'launch', 'load', 's', 'l' then check and start
        when 'stop', 'unload', 'terminate', 'term', 't', 'u' then check and stop
        else
          argh "Unknown command `#{@cmd}`", false
          usage(1)
      end      
    end
    
    # Check if formula has been found
    def check
      argh("Formula missing, please provide a formula name") unless service
      true
    end

    # List all running services with PID and status and path to plist file, if available
    def list
      eww("No %s services controlled by `#{bin}` running..." % [root? ? 'root' : 'user-space']) and return if running.empty?
      running.each do |label|
        if svc = Service.from(label)
          status = !svc.dest.file? ? "#{Tty.red}stale  " : "#{Tty.white}started"
          puts "%-10.10s %s#{Tty.reset}  %7s %s" % [svc.name, status, svc.pid ? svc.pid.to_s : '-', svc.dest.file? ? svc.dest : label]
        else
          puts "%-10.10s #{Tty.red}unknown#{Tty.reset} %7s #{label}" % ["?", "-"]        
        end
      end
    end

    # Kill services without plist file and remove unused plists
    def cleanup
      cleaned = []
      
      # 1. kill services which have no plist file
      running.each do |label|
        if svc = Service.from(label)
          if !svc.dest.file?
            puts "%-15.15s #{Tty.white}stale#{Tty.reset} => killing service..." % svc.name
            kill(svc)
            cleaned << label            
          end
        else
          eww "Service #{label} not managed by `#{bin}` => skipping"
        end
      end
      
      # 2. remove unused plist files
      Dir[path + 'com.github.homebrew.*.plist'].each do |file|
        unless running.include?(File.basename(file).sub(/\.plist$/i, ''))
          puts "Removing unused plist #{file}"
          rm file
          cleaned << file
        end
      end
      
      puts "All #{root? ? 'root' : 'user-space'} services OK, nothing cleaned..." if cleaned.empty?
    end

    # Stop if loaded, then start again
    def restart
      stop if service.loaded?
      start
    end

    # Start a service
    def start
      argh "Service `#{service.name}` already started, use `#{bin} restart #{service.name}`" if service.loaded?
      
      custom_plist = @args.first
      if custom_plist
        if custom_plist =~ %r{\Ahttps?://.+}
          custom_plist = { :url => custom_plist }
        elsif File.exist?(custom_plist)
          custom_plist = Pathname.new(custom_plist)
        else
          argh "#{custom_plist} is not a url or exising file"
        end
      end

      argh "Formula `#{service.name}` not installed, #startup_plist not implemented or no plist file found" if !custom_plist && !service.plist?
      
      temp = Tempfile.new(service.label)
      temp << service.generate_plist(custom_plist)
      temp.flush
  
      rm service.dest if service.dest.exist?
      cp temp.path, service.dest
      
      # clear tempfile
      temp.close
  
      safe_system launchctl, "load", "-w", service.dest.to_s
      $?.to_i != 0 ? argh("Failed to start `#{service.name}`") : ohai("Successfully started `#{service.name}` as #{service.label}")
    end

    # Stop a service or kill if no plist file available...
    def stop
      unless service.loaded?
        rm service.dest if service.dest.exist? # get rid of installed plist anyway, dude
        argh "Service `#{service.name}` not running, wanna start it? Try `#{bin} start #{service.name}`"
      end
  
      if service.dest.exist?
        puts "Stopping `#{service.name}`... (might take a while)"
        safe_system launchctl, "unload", "-w", service.dest.to_s
        $?.to_i != 0 ? argh("Failed to stop `#{service.name}`") : ohai("Successfully stopped `#{service.name}` via #{service.label}")
      else
        puts "Stopping stale service `#{service.name}`... (might take a while)"
        kill(service)
      end
      rm service.dest if service.dest.exist?    
    end
    
    # Kill service without plist file by issuing a `launchctl remove` command
    def kill(svc)
      safe_system launchctl, "remove", svc.label
      argh("Failed to remove `#{svc.name}`, try again?") unless $?.to_i == 0
      while svc.loaded?
        puts "  ...checking status"
        sleep(5)
      end
      ohai "Successfully stopped `#{svc.name}` via #{svc.label}"
    end
  end
end

# Wrapper for a formula to handle service related stuff like parsing
# and generating the plist file.
class Service

  # Support module which will be used to extend Formula with a method :)
  module PlistSupport
    # As a replacement value for `<key>UserName</key>`.
    def startup_user; ServicesCli.user end
  end

  # Access the `Formula` instance
  attr_reader :formula

  # Create a new `Service` instance from either a path or label.
  def self.from(path_or_label)
    return nil unless path_or_label =~ /com\.github\.homebrew\.([^\.]+)(\.plist)?\z/
    new(Formula.factory(Formula.resolve_alias($1))) rescue nil
  end

  # Initialize new `Service` instance with supplied formula.
  def initialize(formula); @formula = formula end
  
  # Delegate access to `formula.name`.
  def name; @name ||= formula.name end

  # Label, static, always looks like `com.github.homebrew.<formula>`.
  def label; @label ||= "com.github.homebrew.#{name}" end

  # Path to a static plist file, this is always `com.github.homebrew.<formula>.plist`.
  def plist; @plist ||= formula.prefix + "#{label}.plist" end

  # Path to destination plist, if run as root it's in `boot_path`, else `user_path`.
  def dest; (ServicesCli.root? ? ServicesCli.boot_path : ServicesCli.user_path) + "#{label}.plist" end

  # Returns `true` if formula implements #startup_plist or file exists.
  def plist?; formula.installed? && (plist.file? || formula.respond_to?(:startup_plist)) end

  # Returns `true` if service is loaded, else false.
  def loaded?; %x{#{ServicesCli.launchctl} list | grep #{label} 2>/dev/null}.chomp =~ /#{label}\z/ end

  # Get current PID of daemon process from launchctl.
  def pid
    status = %x{#{ServicesCli.launchctl} list | grep #{label} 2>/dev/null}.chomp
    return $1.to_i if status =~ /\A([\d]+)\s+.+#{label}\z/
  end

  # Generate that plist file, dude.
  def generate_plist(data = nil)
    data ||= plist.file? ? plist : formula.startup_plist
    
    if data.respond_to?(:file?) && data.file?
      data = data.read      
    elsif data.respond_to?(:keys) && data.keys.include?(:url)
      require 'open-uri'
      data = open(data).read
    end
  
    # replace "template" variables and ensure label is always, always com.github.homebrew.<formula>
    data = data.to_s.gsub(/\{\{([a-z][a-z0-9_]*)\}\}/i) { |m| formula.send($1).to_s if formula.respond_to?($1) }.
              gsub(%r{(<key>Label</key>\s*<string>)[^<]*(</string>)}, '\1' + label + '\2')

    # and force fix UserName, if necessary
    if formula.startup_user != "root" && data =~ %r{<key>UserName</key>\s*<string>root</string>}    
      data = data.gsub(%r{(<key>UserName</key>\s*<string>)[^<]*(</string>)}, '\1' + formula.startup_user + '\2')
    elsif ServicesCli.root? && formula.startup_user != "root" && data !~ %r{<key>UserName</key>}
      data = data.gsub(%r{(</dict>\s*</plist>)}, "  <key>UserName</key><string>#{formula.startup_user}</string>\n\\1")
    end

    if ARGV.verbose?
      ohai "Generated plist for #{formula.name}:"
      puts "   " + data.gsub("\n", "\n   ")
      puts
    end

    data
  end
end

# Start the cli dispatch stuff.
#
ServicesCli.run!