openwferu-extras 0.9.12.863

data/lib/openwfe/extras/util/sqs.rb

@@ -0,0 +1,581 @@
+#
+#--
+# Copyright (c) 2007, John Mettraux, OpenWFE.org
+# All rights reserved.
+# 
+# Redistribution and use in source and binary forms, with or without 
+# modification, are permitted provided that the following conditions are met:
+# 
+# . Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.  
+# 
+# . Redistributions in binary form must reproduce the above copyright notice, 
+#   this list of conditions and the following disclaimer in the documentation 
+#   and/or other materials provided with the distribution.
+# 
+# . Neither the name of the "OpenWFE" nor the names of its contributors may be
+#   used to endorse or promote products derived from this software without
+#   specific prior written permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
+# POSSIBILITY OF SUCH DAMAGE.
+#++
+#
+# $Id$
+#
+
+#
+# Made in Japan
+#
+# John dot Mettraux at OpenWFE dot org
+#
+
+require 'base64'
+require 'cgi'
+require 'net/https'
+require 'rexml/document'
+require 'time'
+require 'pp'
+
+
+module SQS
+
+    #
+    # An SQS message (after its creation).
+    #
+    class Message
+
+        attr_reader :queue, :message_id, :message_body
+
+        def initialize (queue, xml_element)
+
+            @queue = queue
+            @message_id = SQS::get_element_text(xml_element, "MessageId")
+            @message_body = SQS::get_element_text(xml_element, "MessageBody")
+        end
+
+        #
+        # Connects to the queue service and deletes this message in its queue.
+        #
+        def delete
+            @queue.queue_service.delete_message(@queue, @message_id)
+        end
+    end
+
+    #
+    # An SQS queue (gathering all the necessary info about it 
+    # in a single class).
+    #
+    class Queue
+
+        attr_reader :queue_service, :host, :path, :name
+
+        def initialize (queue_service, xml_element)
+
+            @queue_service = queue_service
+
+            s = xml_element.text.to_s
+            m = Regexp.compile('^http://(.*)(/.*)(/.*$)').match(s)
+            @host = m[1]
+            @name = m[3][1..-1]
+            @path = m[2] + m[3]
+        end
+    end
+
+    class QueueService
+
+        AWS_VERSION = "2006-04-01"
+        DEFAULT_QUEUE_HOST = "queue.amazonaws.com"
+
+        def initialize (queue_host=nil)
+
+            @queue_host = queue_host
+            @queue_host = DEFAULT_QUEUE_HOST unless @queue_host
+        end
+
+        #
+        # Lists the queues for the active AWS account.
+        # If 'prefix' is given, only queues whose name begin with that
+        # prefix will be returned.
+        #
+        def list_queues (prefix=nil)
+
+            queues = []
+
+            path = "/"
+            path = "#{path}?QueueNamePrefix=#{prefix}" if prefix
+
+            doc = do_action :get, @queue_host, path
+
+            doc.elements.each("//QueueUrl") do |e|
+                queues << Queue.new(self, e)
+            end
+
+            return queues
+        end
+
+        #
+        # Creates a queue.
+        #
+        # If the queue name doesn't comply with SQS requirements for it,
+        # an error will be raised.
+        #
+        def create_queue (queue_name)
+
+            doc = do_action :post, @queue_host, "/?QueueName=#{queue_name}"
+
+            doc.elements.each("//QueueUrl") do |e|
+                return e.text.to_s
+            end
+        end
+
+        #
+        # Given some content ('text/plain' content), send it as a message to
+        # a queue.
+        # Returns the SQS message id (a String).
+        #
+        # The queue might be a queue name (String) or a Queue instance.
+        #
+        def put_message (queue, content)
+
+            queue = resolve_queue(queue)
+
+            doc = do_action :put, queue.host, "#{queue.path}/back", content
+
+            #puts doc.to_s
+
+            #status_code = SQS::get_element_text(doc, '//StatusCode')
+            #message_id = SQS::get_element_text(doc, '//MessageId')
+            #request_id = SQS::get_element_text(doc, '//RequestId')
+            #{ :status_code => status_code, 
+            #  :message_id => message_id, 
+            #  :request_id => request_id }
+
+            SQS::get_element_text(doc, '//MessageId')
+        end
+
+        alias :send_message :put_message
+
+        #
+        # Retrieves a bunch of messages from a queue. Returns a list of
+        # Message instances.
+        #
+        # There are actually two optional params that this method understands :
+        #
+        # - :timeout  the duration in seconds of the message visibility in the
+        #             queue
+        # - :count    the max number of message to be returned by this call
+        #
+        # The queue might be a queue name (String) or a Queue instance.
+        #
+        def get_messages (queue, params={})
+
+            queue = resolve_queue(queue)
+
+            path = "#{queue.path}/front"
+
+            path += "?" if params.size > 0
+
+            timeout = params[:timeout]
+            count = params[:count]
+
+            path += "VisibilityTimeout=#{timeout}" if timeout
+            path += "&" if timeout and count
+            path += "NumberOfMessages=#{count}" if count
+
+            doc = do_action :get, queue.host, path
+
+            messages = []
+
+            doc.elements.each("//Message") do |me|
+                messages << Message.new(queue, me)
+            end
+
+            messages
+        end
+
+        #
+        # Retrieves a single message from a queue. Returns an instance of
+        # Message.
+        #
+        # The queue might be a queue name (String) or a Queue instance.
+        #
+        def get_message (queue, message_id)
+
+            queue = resolve_queue(queue)
+
+            path = "#{queue.path}/#{message_id}"
+
+            begin
+                doc = do_action :get, queue.host, path
+                Message.new(queue, doc.root.elements[1])
+            rescue Exception => e
+                #puts e.message
+                return nil if e.message.match "^404 .*$"
+                raise e
+            end
+        end
+
+        #
+        # Deletes a given message.
+        #
+        # The queue might be a queue name (String) or a Queue instance.
+        #
+        def delete_message (queue, message_id)
+
+            queue = resolve_queue(queue)
+
+            path = "#{queue.path}/#{message_id}"
+            #path = "#{queue.path}/#{CGI::escape(message_id)}"
+
+            doc = do_action :delete, queue.host, path
+
+            SQS::get_element_text(doc, "//StatusCode") == "Success"
+        end
+
+        #
+        # Use with care !
+        #
+        # Attempts at deleting all the messages in a queue.
+        # Returns the total count of messages deleted.
+        #
+        # A call on this method might take a certain time, as it has
+        # to delete each message individually. AWS will perhaps
+        # add a proper 'flush_queue' method later.
+        #
+        # The queue might be a queue name (String) or a Queue instance.
+        #
+        def flush_queue (queue)
+
+            count = 0
+
+            while true
+
+                l = get_messages(queue, :timeout => 0, :count => 255)
+                break if l.length < 1
+
+                l.each do |m|
+                    m.delete
+                    count += 1
+                end
+            end
+
+            return count
+        end
+
+        #
+        # Deletes the queue. Returns true if the delete was successful.
+        # You can empty a queue by called the method #flush_queue
+        #
+        # If 'force' is set to true, a flush will be performed on the
+        # queue before the actual delete operation. It should ensure
+        # a successful removal of the queue.
+        #
+        def delete_queue (queue, force=false)
+
+            queue = resolve_queue(queue)
+
+            flush_queue(queue) if force
+            
+            begin
+
+                doc = do_action :delete, @queue_host, queue.path
+
+            rescue Exception => e
+
+                return false if e.message.match "^400 .*$"
+            end
+
+            SQS::get_element_text(doc, "//StatusCode") == "Success"
+        end
+
+        #
+        # Given a queue name, a Queue instance is returned.
+        #
+        def get_queue (queue_name)
+
+            l = list_queues(queue_name)
+
+            l.each do |q|
+                return q if q.name == queue_name
+            end
+
+            #return nil
+            raise "found no queue named '#{queue_name}'"
+        end
+
+        protected
+
+            #
+            # 'queue' might be a Queue instance or a queue name. 
+            # If it's a Queue instance, it is immediately returned,
+            # else the Queue instance is looked up and returned.
+            #
+            def resolve_queue (queue)
+                return queue if queue.kind_of? Queue
+                return get_queue(queue.to_s)
+            end
+
+            #
+            # The actual http request/response job is done here.
+            #
+            def do_action (action, host, path, content=nil)
+
+                #puts "___path : #{path}"
+
+                doc = nil
+
+                http = Net::HTTP.new(host)
+                http.start do
+
+                    date = Time.now.httpdate
+
+                    req = if action == :get
+                        Net::HTTP::Get.new(path)
+                    elsif action == :post
+                        Net::HTTP::Post.new(path)
+                    elsif action == :put
+                        Net::HTTP::Put.new(path)
+                    else #action == :delete
+                        Net::HTTP::Delete.new(path)
+                    end
+
+                    req['AWS-Version'] = AWS_VERSION
+                    req['Date'] = date
+                    req['Content-type'] = 'text/plain'
+
+                    if action == :put or action == :post
+                        req.body = content
+                        req['Content-length'] = content.length.to_s if content
+                    end
+
+                    req['Authorization'] = generate_auth_header(
+                        action, path, date, "text/plain")
+
+                    #req.each_header do |k, v|
+                    #    puts " - '#{k}' => '#{v}'"
+                    #end
+
+                    res = http.request(req)
+
+                    case res
+                    when Net::HTTPSuccess, Net::HTTPRedirection
+                        doc = REXML::Document.new(res.read_body)
+                    else
+                        res.error!
+                    end
+                end
+                raise_errors(doc)
+                return doc
+            end
+
+            #
+            # Scans the SQS XML reply for potential errors and raises an
+            # error if he encounters one.
+            #
+            def raise_errors (doc)
+
+                doc.elements.each("//Error") do |e|
+
+                    code = get_element_text(e, "Code")
+                    return unless code
+
+                    message = get_element_text(e, "Message")
+                    raise "SQS::#{code} : #{m.text.to_s}"
+                end
+            end
+
+            #
+            # Generates the 'AWS x:y" authorization header value.
+            #
+            def generate_auth_header (action, path, date, content_type)
+
+                s = ""
+                s << action.to_s.upcase
+                s << "\n"
+
+                #s << Base64.encode64(Digest::MD5.digest(content)).strip \
+                #    if content
+                    #
+                    # documented but not necessary (not working)
+                s << "\n"
+
+                s << content_type
+                s << "\n"
+
+                s << date
+                s << "\n"
+
+                i = path.index '?'
+                path = path[0..i-1] if i
+                s << path
+
+                #puts ">>>#{s}<<<"
+
+                digest = OpenSSL::Digest::Digest.new 'sha1'
+
+                key = ENV['AMAZON_SECRET_ACCESS_KEY']
+
+                raise "No $AMAZON_SECRET_ACCESS_KEY env variable found" \
+                    unless key
+
+                sig = OpenSSL::HMAC.digest(digest, key, s)
+                sig = Base64.encode64(sig).strip
+
+                "AWS #{ENV['AMAZON_ACCESS_KEY_ID']}:#{sig}"
+            end
+
+    end
+
+    #
+    # A convenience method for returning the text of a sub element,
+    # maybe there is something better in REXML, but I haven't found out
+    # yet.
+    #
+    def SQS.get_element_text (parent_elt, elt_name)
+        e = parent_elt.elements[elt_name]
+        return nil unless e
+        return e.text.to_s
+    end
+end
+
+
+#
+# running directly...
+
+if $0 == __FILE__
+
+    if ENV['AMAZON_ACCESS_KEY_ID'] == nil or ENV['AMAZON_SECRET_ACCESS_KEY'] == nil
+
+        puts
+        puts "env variables $AMAZON_ACCESS_KEY_ID and $AMAZON_SECRET_ACCESS_KEY are not set"
+        puts
+        exit 1
+    end
+
+    ACTIONS = {
+        :list_queues => :list_queues,
+        :lq => :list_queues,
+        :create_queue => :create_queue,
+        :cq => :create_queue,
+        :delete_queue => :delete_queue,
+        :dq => :delete_queue,
+        :flush_queue => :flush_queue,
+        :fq => :flush_queue,
+        :get_message => :get_message,
+        :gm => :get_message,
+        :delete_message => :delete_message,
+        :dm => :delete_message,
+        :puts_message => :put_message,
+        :pm => :put_message
+    }
+
+    b64 = false
+    queue_host = nil
+
+    require 'optparse'
+
+    opts = OptionParser.new
+
+    opts.banner = "Usage: sqs.rb [options] {action} [queue_name] [message_id]"
+    opts.separator("")
+    opts.separator("   known actions are :")
+    opts.separator("")
+
+    keys = ACTIONS.keys.collect { |k| k.to_s }.sort
+    keys.each { |k| opts.separator("      - '#{k}'  (#{ACTIONS[k.intern]})") }
+
+    opts.separator("")
+    opts.separator("   options are :")
+    opts.separator("")
+
+    opts.on("-H", "--host", "AWS queue host") do |host|
+        queue_host = host
+    end
+
+    opts.on("-h", "--help", "displays this help / usage") do
+        STDERR.puts "\n#{opts.to_s}\n"
+        exit 0
+    end
+
+    opts.on("-b", "--base64", "encode/decode messages with base64") do
+        b64 = true
+    end
+
+    argv = opts.parse(ARGV)
+
+    if argv.length < 1
+        STDERR.puts "\n#{opts.to_s}\n"
+        exit 0
+    end
+
+    a = argv[0]
+    queue_name = argv[1]
+    message_id = argv[2]
+
+    action = ACTIONS[a.intern]
+
+    unless action
+        STDERR.puts "unknown action '#{a}'"
+        exit 1
+    end
+
+    qs = SQS::QueueService.new
+
+    STDERR.puts "#{action.to_s}..."
+
+    #
+    # just do it
+
+    case action
+    when :list_queues, :create_queue, :delete_queue, :flush_queue
+
+        pp qs.send(action, queue_name)
+
+    when :get_message
+
+        if message_id
+            m = qs.get_message(queue_name, message_id)
+            body = m.message_body
+            body = Base64.decode64(body) if b64
+            puts body
+        else
+            pp qs.get_messages(queue_name, :timeout => 0, :count => 255)
+        end
+
+    when :delete_message
+
+        raise "argument 'message_id' is missing" unless message_id
+        pp qs.delete_message(queue_name, message_id)
+
+    when :put_message
+
+        message = argv[2]
+
+        unless message
+            message = ""
+            while true
+                s = STDIN.gets()
+                break if s == nil
+                message += s[0..-2]
+            end
+        end
+
+        message = Base64.encode64(message).strip if b64
+
+        pp qs.put_message(queue_name, message)
+    else
+
+        STDERR.puts "not yet implemented..."
+    end
+
+end