copied over download.rb changes from oldbooru

app/logical/download.rb

@@ -1,27 +1,26 @@
 class Download
   class Error < Exception ; end
   
-  attr_accessor :source, :content_type
+  attr_accessor :source, :content_type, :file_path
   
   def initialize(source, file_path)
     @source = source
     @file_path = file_path
   end
   
-  # Downloads to @file_path
   def download!
-    http_get_streaming(@source) do |response|
+    http_get_streaming do |response|
       self.content_type = response["Content-Type"]
-      File.open(@file_path, "wb") do |out|
+      File.open(file_path, "wb") do |out|
         response.read_body(out)
       end
     end
-    @source = fix_image_board_sources(@source)
+    after_download
   end
   
-  # private
+  def pixiv_rewrite(headers)
-  def handle_pixiv(source, headers)
+    return unless source =~ /pixiv\.net/
-    if source =~ /pixiv\.net/
+
     headers["Referer"] = "http://www.pixiv.net"
     
     # Don't download the small version
@@ -29,12 +28,42 @@ class Download
       match = $1
       source.sub!(match + "_m", match)
     end
+    
+    # Download the big version if it exists
+    if source =~ %r!(\d+_p\d+)\.!
+      match = $1
+      repl = match.sub(/_p/, "_big_p")
+      big_source = source.sub(match, repl)
+      if pixiv_http_exists?(big_source)
+        self.source = big_source
+      end
+    end
   end
   
-    source
+  def pixiv_http_exists?
+    # example: http://img01.pixiv.net/img/as-special/15649262_big_p2.jpg
+    exists = false
+    uri = URI.parse(source)
+    Net::HTTP.start(uri.host, uri.port) do |http|
+      headers = {"Referer" => "http://www.pixiv.net", "User-Agent" => "#{Danbooru.config.app_name}/#{Danbooru.config.version}"}
+      http.request_head(uri.request_uri, headers) do |res|
+        if res.is_a?(Net::HTTPSuccess)
+          exists = true
+        end
+      end
+    end
+    exists
   end
   
-  def http_get_streaming(source, options = {})
+  def before_download(headers)
+    pixiv_rewrite(headers)
+  end
+  
+  def after_download
+    fix_image_board_sources
+  end
+  
+  def http_get_streaming(options = {})
     max_size = options[:max_size] || Danbooru.config.max_file_size
     max_size = nil if max_size == 0 # unlimited
     limit = 4
@@ -51,7 +80,7 @@ class Download
         headers = {
           "User-Agent" => "#{Danbooru.config.safe_app_name}/#{Danbooru.config.version}"
         }
-        source = handle_pixiv(source, headers)
+        before_download(headers)
         url = URI.parse(source)
         http.request_get(url.request_uri, headers) do |res|
           case res
@@ -78,11 +107,9 @@ class Download
     end # while
   end # def
   
-  def fix_image_board_sources(source)
+  def fix_image_board_sources
     if source =~ /\/src\/\d{12,}|urnc\.yi\.org|yui\.cynthia\.bne\.jp/
-      "Image board"
+      self.source = "Image board"
-    else
-      source
     end
   end
 end

test/unit/download_test.rb

@@ -13,7 +13,7 @@ class DownloadTest < ActiveSupport::TestCase
     end
     
     should "stream a file from an HTTP source" do
-      @download.http_get_streaming(@download.source) do |resp|
+      @download.http_get_streaming do |resp|
         assert_equal("200", resp.code)
         assert(resp["Content-Length"].to_i > 0, "File should be larger than 0 bytes")
       end
@@ -21,7 +21,7 @@ class DownloadTest < ActiveSupport::TestCase
     
     should "throw an exception when the file is larger than the maximum" do
       assert_raise(Download::Error) do
-        @download.http_get_streaming(@download.source, :max_size => 1) do |resp|
+        @download.http_get_streaming(:max_size => 1) do |resp|
         end
       end
     end