implement iqdb wrapper

app/logical/iqdb/download.rb

@@ -0,0 +1,21 @@
+module Iqdb
+  class Download
+    attr_reader :source, :download, :similar_posts
+
+    def initialize(source)
+      @source = source
+    end
+
+    def download_from_source
+      tempfile = Tempfile.new("iqdb-#{$PROCESS_ID}")
+      @download = Downloads::File.new(source, tempfile.path)
+      @download.download!
+    end
+
+    def find_similar
+      if Danbooru.config.iqdb_hostname_and_port
+        @similar_posts = Iqdb::Server.new(*Danbooru.config.iqdb_hostname_and_port).query(0, 3, @download.file_path)
+      end
+    end
+  end
+end

app/logical/iqdb/importer.rb

@@ -0,0 +1,13 @@
+module Iqdb
+  class Importer
+    def import!
+      IO.popen("iqdb add #{Rails.root}/danbooru.db", "w") do |io|
+        Post.find_each do |post|
+          if File.exists?(post.preview_file_path)
+            io.puts "#{post.id.to_s(16)}:#{post.preview_file_path}"
+          end
+        end
+      end
+    end
+  end
+end

app/logical/iqdb/responses/base.rb

@@ -0,0 +1,6 @@
+module Iqdb
+  module Responses
+    class Base
+    end
+  end
+end

app/logical/iqdb/responses/collection.rb

@@ -0,0 +1,27 @@
+module Iqdb
+  module Responses
+    class Collection
+      def initialize(response_string)
+        @responses = response_string.split(/\n/).map do |string|
+          ::Iqdb::Responses.const_get("Response_#{string[0..2]}").new(string[4..-1])
+        end
+      end
+
+      def matches
+        @matches ||= responses.select {|x| x.is_a?(Iqdb::Responses::Response_200) && x.score >= 0.9}
+      end
+
+      def empty?
+        matches.empty?
+      end
+
+      def errored?
+        errors.any?
+      end
+
+      def errors
+        @errors ||= responses.select {|x| x.is_a?(Iqdb::Responses::Error)}.map {|x| x.to_s}
+      end
+    end
+  end
+end

app/logical/iqdb/responses/error.rb

@@ -0,0 +1,6 @@
+module Iqdb
+  module Responses
+    class Error < Base
+    end
+  end
+end

app/logical/iqdb/responses/response_000.rb

@@ -0,0 +1,8 @@
+module Iqdb
+  module Responses
+    class Response_000 < Base
+      def initialize(response_string)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_100.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_100 < Base
+      attr_reader :message
+
+      def initialize(response_string)
+        @message = response_string
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_101.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_101 < Base
+      attr_reader :key, :value
+
+      def initialize(response_string)
+        @key, @value = response_string.split(/\=/)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_102.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_102 < Base
+      attr_reader :dbid, :filename
+
+      def initialize(response_string)
+        @dbid, @filename = response_string.split(/ /)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_200.rb

@@ -0,0 +1,18 @@
+module Iqdb
+  module Responses
+    class Response_200 < Base
+      attr_reader :imgid, :score, :width, :height
+
+      def initialize(response_string)
+        @imgid, @score, @width, @height = response_string.split(/ /)
+        @score = score.to_f
+        @width = width.to_i
+        @height = height.to_i
+      end
+
+      def post_id
+        imgid.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_201.rb

@@ -0,0 +1,19 @@
+module Iqdb
+  module Responses
+    class Response_201 < Base
+      attr_reader :dbid, :imgid, :score, :width, :height
+
+      def initialize(response_string)
+        @dbid, @imgid, @score, @width, @height = response_string.split(/ /)
+        @dbid = dbid.to_i
+        @score = score.to_f
+        @width = width.to_i
+        @height = height.to_i
+      end
+
+      def post_id
+        imgid.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_202.rb

@@ -0,0 +1,19 @@
+module Iqdb
+  module Responses
+    class Response_202 < Base
+      attr_reader :original_id, :stddev, :dupes
+
+      def initialize(response_string)
+        response_string =~ /^(\d+)=([0-9.]+)/
+        @original_id = $1
+        @stddev = $2
+
+        @dupes = response_string.scan(/(\d+):([0-9.]+)/).map {|x| [x[0].to_i(16), x[1].to_f]}
+      end
+
+      def original_post_id
+        original_id.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_300.rb

@@ -0,0 +1,15 @@
+module Iqdb
+  module Responses
+    class Response_300 < Error
+      attr_reader :message
+
+      def initialize(response_string)
+        @message = response_string
+      end
+
+      def to_s
+        "Error: #{message}"
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_301.rb

@@ -0,0 +1,17 @@
+module Iqdb
+  module Responses
+    class Response_301 < Error
+      attr_reader :exception, :description
+
+      def initialize(response_string)
+        response_string =~ /^(\S+) (.+)/
+        @exception = $1
+        @description = $2
+      end
+
+      def to_s
+        "Exception: #{exception}: #{description}"
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_302.rb

@@ -0,0 +1,17 @@
+module Iqdb
+  module Responses
+    class Response_302 < Error
+      attr_reader :exception, :description
+
+      def initialize(response_string)
+        response_string =~ /^(\S+) (.+)/
+        @exception = $1
+        @description = $2
+      end
+
+      def to_s
+        "Fatal Exception: #{exception}: #{description}"
+      end
+    end
+  end
+end

app/logical/iqdb/server.rb

@@ -0,0 +1,56 @@
+module Iqdb
+  class Server
+    FLAG_SKETCH = 0x01
+    FLAG_GRAYSCALE = 0x02
+    FLAG_WIDTH_AS_SET = 0x08
+    FLAG_DISCARD_COMMON_COEFFS = 0x16
+
+    attr_reader :hostname, :port
+
+    def self.import(database)
+      IO.popen("iqdb #{database}", "w") do |io|
+        Post.find_each do |post|
+          puts "Adding #{post.id}"
+          io.puts "#{post.id.to_s(16)} :#{post.preview_file_path}"
+        end
+      end
+    end
+
+    def self.add(database, image_id, filename)
+      image_id_hex = image_id.to_s(16)
+      `iqdb add #{database} #{image_id_hex} :#{filename}`
+    end
+
+    def self.remove(database, image_id)
+      image_id_hex = image_id.to_s(16)
+      `iqdb remove 0 #{image_id_hex} #{database}`
+    end
+
+    def initialize(hostname, port)
+      @hostname = hostname
+      @port = port
+    end
+
+    def open
+      @socket = TCPSocket.new(hostname, port)
+    end
+
+    def close
+      @socket.close
+    end
+
+    def request
+      open
+      yield
+    ensure
+      close
+    end
+
+    def query(dbid, results, filename, flags = FLAG_DISCARD_COMMON_COEFFS)
+      request do
+        @socket.puts "query #{dbid} #{flags} #{results} #{filename}"
+        responses = Responses::Collection.new(@socket.read)
+      end
+    end
+  end
+end