ニコニコ動画のダウンロードツールを改良してみた。（その2）

昨日の続き。

今日は、アプリケーション側のコードを実装していく。

中断ファイルの処理

中断ファイルから中断時の情報を取り出す必要があるので、それに関するコードをまず用意。

最初に、中断情報を表す構造体を定義しておく。

module NicovideoDL
  module LibRTMP
    ResumeInfo = Struct.new(:meta_header, :meta_header_size,
                            :last_key_frame, :last_key_frame_type, :last_key_frame_size,
                            :seek, :duration, :last_file_position)
  end
end

これで、ResumeInfoというクラスと、付随するアクセサが定義される。

あと、この中断情報を使って、ファイルがほぼダウンロード済みかという終了判定を行う必要があるので、次のメソッドを定義しておく。

module NicovideoDL
  module LibRTMP
    class ResumeInfo
      # Workaround to exit with 0 if the file is fully (> 99.9%) downloaded
      def fully_downloaded?
        seek = self.seek
        duration = self.duration
        (duration > 0) && (seek >= (duration * 999))
      end
    end
  end
end

次に、この中断情報をFLVファイルから取得できるようにするために、Fileクラスを継承したFLVFileクラスを用意し、get_resume_infoというメソッドを実装する。
やってる内容は、OpenResumeFile()とGetLastKeyframe()を合わせたもの。

module NicovideoDL
  module LibRTMP
    class FLVFile < File
      def get_resume_info
        resume_info = ResumeInfo.new

        file_size = self.size

        # read file header

        self.pos = 0
        header = self.read(13)
        if (header.byteslice(0) != 'F') ||
            (header.byteslice(1) != 'L') ||
            (header.byteslice(2) != 'V') ||
            (header.getbyte(3) != 0x01)
          raise "invalid FLV file."
        end

        file_type = header.getbyte(4)
        if (file_type & 0x05) == 0
          raise "FLV file contains neither video nor audio."
        end
        audio_only = ((file_type & 0x04) > 0) && ((file_type & 0x01) == 0x00)

        data_offset = LibRTMP.decode_int32(header.byteslice(5, 4))
        self.pos = data_offset

        prev_tag_size_data = self.read(4)
        prev_tag_size = LibRTMP.decode_int32(prev_tag_size_data)
        if prev_tag_size != 0
          raise "first prev tag size is not 0."
        end

        # find meta data

        block_offset = data_offset + 4
        found = false
        while (block_offset < (file_size - 4)) && (!found)
          self.pos = block_offset
          block_header = self.read(4)
          data_size = LibRTMP.decode_int24(block_header.byteslice(1, 3))
          if block_header.getbyte(0) == 0x12
            self.pos = block_offset + 11
            data = self.read(data_size)
            data_buf = FFI::MemoryPointer.new :char, data_size
            data_buf.write_bytes(data, 0, data_size)
            meta_obj = AMFObject.new
            if meta_obj.decode(data_buf, data_size, false) < 0
              raise "failed to decode meta data packet."
            end
            property = meta_obj.get_property(nil, 0)
            meta_string = property.get_string

            if meta_string == "onMetaData"
              meta_obj.dump
              resume_info.meta_header_size = data_size
              resume_info.meta_header = data_buf
              prop_duration = meta_obj.find_first_matching_property("duration")
              if prop_duration
                resume_info.duration = prop_duration.get_number
              else
                resume_info.duration = 0
              end
              found = true
              break
            end
          end
          block_offset += 11 + data_size + 4
        end

        unless found
          raise "failed to find meta data."
        end

        # find last key frame

        block_offset = file_size
        block_header = nil
        prev_tag_size = 0
        loop do
          if block_offset < 13
            raise "unexpected start of file."
          end

          self.pos = block_offset - 4
          prev_tag_size_data = self.read(4)
          prev_tag_size = LibRTMP.decode_int32(prev_tag_size_data)
          if prev_tag_size == 0
            raise "faild to find last key frame."
          end
          if (prev_tag_size < 0) || (prev_tag_size > (block_offset - 4 - 13))
            raise "invalid prev tag size."
          end

          block_offset -= (prev_tag_size + 4)
          self.pos = block_offset
          block_header = self.read(11)
          data_start = self.read(1)

          if audio_only && (block_header.getbyte(0) == 0x08)
            break
          end
          if (!audio_only) && (block_header.getbyte(0) == 0x09) && ((data_start.getbyte(0) & 0xf0) == 0x10)
            break
          end
        end

        resume_info.last_key_frame_type = block_header.getbyte(0)
        last_key_frame_size = prev_tag_size - 11
        resume_info.last_key_frame_size = last_key_frame_size
        self.pos = block_offset + 11
        last_key_frame_data = self.read(last_key_frame_size)
        last_key_frame = FFI::MemoryPointer.new :char, last_key_frame_size
        last_key_frame.write_bytes(last_key_frame_data, 0, last_key_frame_size)
        resume_info.last_key_frame = last_key_frame

        resume_info.seek = LibRTMP.decode_int24(block_header.byteslice(4, 3))
        resume_info.seek |= (block_header.getbyte(7) << 24)

        resume_info.last_file_position = block_offset + prev_tag_size + 4

        return resume_info
      end
    end
  end
end

これで中断ファイルの処理はOK。

RTMPの接続、データの取得

まず、RTMPのURLをパースする必要があるので、次のようなファクトリメソッドを用意。
（RTMP自体はFFI::Structを継承しているので、initializeをオーバーライトしてしまうのはちょっと危険かなと判断）

module NicovideoDL
  module LibRTMP
    class RTMP
      def self.create_from_uri(uri, page_uri)
        rtmp = self.new
        rtmp.clear
        rtmp.init
        rtmp.setup_uri(uri, page_uri)
        rtmp
      end

      def setup_uri(uri, page_uri)
        url_without_query = URI::Generic.build(scheme: uri.scheme, host: uri.host, path: uri.path).to_s
        @tc_url = url_without_query
        @page_url = page_uri.to_s
        @playpath = uri.query.split('=')[1]

        url_info = LibRTMP.parse_url(url_without_query)

        @protocol = url_info[:protocol]
        @host = url_info[:host]
        @port = url_info[:port]
        @playpath = url_info[:playpath] unless url_info[:playpath].empty?
        @app = url_info[:app]

        if @port == 0
          if (@port & RTMP_FEATURE_SSL) > 0
            @port = 443
          elsif (@port & RTMP_FEATURE_HTTP) > 0
            @port = 80
          else
            @port = 1935
          end
        end
      end
    end
  end
end

そして、接続する処理を実装。

module NicovideoDL
  module LibRTMP
    class RTMP
      def open(resume_info, &block)
        seek = resume_info ? resume_info.seek : 0

        self.setup_stream(
          @protocol, @host, @port, "",
          @playpath, @tc_url, "",
          @page_url, @app, "",
          "", 0,
          "", "", "",
          0, 0, false, 30)

        self[:Link][:lFlags] |= RTMP_LF_BUFX

        self.set_buffer_ms(DEF_BUFTIME)

        if self.connect(nil) <= 0
          raise "failed to connect."
        end

        if self.connect_stream(seek) <= 0
          raise "failed to connect stream."
        end

        if resume_info
          self[:m_read][:timestamp] = seek
          self[:m_read][:flags] |= RTMP_READ_RESUME
          self[:m_read][:initialFrameType] = resume_info.last_key_frame_type
          self[:m_read][:nResumeTS] = seek
          self[:m_read][:metaHeader] = resume_info.meta_header
          self[:m_read][:nMetaHeaderSize] = resume_info.meta_header_size
          self[:m_read][:initialFrame] = resume_info.last_key_frame
          self[:m_read][:nInitialFrameSize] = resume_info.last_key_frame_size
        end

        block.call

        self.close
      end
    end
  end
end

なお、ホントはconnectというメソッド名にしたかったけど、RTMP_Connect()のラップでconnectというメソッドをすでに定義しているので、メソッド名はopenにしてある。
また、接続したら最後は必ず切断処理をしないといけないので、接続中に行う処理はブロックに委譲するようにしてある。

最後に、データの取得処理を実装。

module NicovideoDL
  module LibRTMP
    class RTMP
      def read_data(&block)
        status = RD_SUCCESS
        buffer_size = 64 * 1024
        buffer = FFI::MemoryPointer.new(:char, buffer_size)
        buffer_time = DEF_BUFTIME

        duration = self.get_duration

        # download

        loop do
          size = self.read(buffer, buffer_size)

          if size > 0
            data = buffer.read_bytes(size)

            if duration <= 0
              duration = self.get_duration
            end
            if (duration > 0) && (buffer_time < (duration * 1000))
              buffer_time = (duration * 1000 + 5000).to_i
              self.set_buffer_ms(buffer_time)
              self.update_buffer_ms
            end

            timestamp = self[:m_read][:timestamp]

            block.call(data, timestamp / 1000.0, duration)
          else
            if self[:m_read][:status] == RTMP_READ_EOF
              break
            elsif self[:m_read][:status] == RTMP_READ_COMPLETE
              break
            end
          end

          if (size < 0) || (!self.connected?) || self.timedout?
            status = RD_INCOMPLETE
            break
          end
        end

        timestamp = self[:m_read][:timestamp]
        if (duration > 0) && (timestamp < (duration * 999))
          status = RD_INCOMPLETE
        end

        status
      end
    end
  end
end

これも、ホントはreadというメソッド名にしたかったけど、RTMP_Read()のラップでreadというメソッドをすでに定義しているので、メソッド名はread_dataに。
そして、得たデータをどう扱うかは、ブロックに委譲するようにしている。

なお、本家のrtmpdumpと違って、再接続の処理は実装してない。
これは、

• 最初は再接続の処理も実装していたんだけど、再接続がうまくいった試しがなかった。（どうやら大元のソケットの接続が切れてしまっているらしく、RTMPオブジェクトを作り直さないとダメ）
• 何回か再接続できたとしても、（rtmpdumpを使ったときと同様に）最後にはどうせRTMPオブジェクトを作り直すことになるのだから、再接続を試すのもあまり意味がない。

ので、だったら最初から再接続の処理は実装せずに、コードをシンプルに保った方がいいと判断したから。

これで、RTMPの接続、データの取得処理も実装できた。
あとは、これらを使うコードだけ。

ダウンロード処理

ダウンロード処理は、以前に実装していたNicovideoDL::Downloaderのプライベートメソッドdownload_with_rtmpを修正。

module NicovideoDL
  class Downloader
     ...
    def download_with_rtmpe(video_info, video_cookies, output_filename)
      original_uri = video_info["original_uri"]
      uri = video_info["uri"]
      playpath = uri.query.split('=')[1]
      fmst1 = video_info["fmst1"]
      fmst2 = video_info["fmst2"]

      status = LibRTMP::RD_SUCCESS
      resume_info = nil
      download_status_printer = DownloadStatusPrinter.new(0, :sec)

      LibRTMP::FLVFile.open(output_filename, "w+b") do |file|
        loop do
          if resume_info && resume_info.fully_downloaded?
            status = LibRTMP::RD_SUCCESS
            break
          end

          rtmp = LibRTMP::RTMP.create_from_uri(uri, original_uri)
          rtmp.add_connect_option("S:#{video_info['fmst1']}")
          rtmp.add_connect_option("S:#{video_info['fmst2']}")
          rtmp.add_connect_option("S:#{playpath}")

          rtmp.open(resume_info) do 
            status = rtmp.read_data do |data, current_sec, total_sec|
              file.write(data)

              if (download_status_printer.total_size == 0) && (total_sec > 0)
                download_status_printer.total_size = total_sec
              end
              download_status_printer.print(current_sec)
            end
          end

          if status == LibRTMP::RD_INCOMPLETE
            resume_info = file.get_resume_info
            file.pos = resume_info.last_file_position
          else
            break
          end
        end
      end
    end
    ...
  end
end

なお、ダウンロードの進捗度合いを表示するためのユーティリティーとして実装したNicovideoDL::DownloadStatusPrinterは、元はデータのサイズ（バイト）で進捗度合いを計算していたけれど、RTMPによるダウンロードの場合、現在のタイムスタンプ（秒）と動画全体の時間の長さ（秒）しか分からないので、単位を「バイト」か「秒」か選べるように修正してある。

具体的には、以下のような実装。

module NicovideoDL
  class DownloadStatusPrinter
    Size1K = 1024
    Epsilon = 0.0001

    def initialize(total_size=0, unit=:byte)
      @start_time = Time.now
      @total_size = total_size
      @unit = unit
    end

    attr_accessor :total_size

    def print(current_size)
      if @total_size != 0
        percent = (100.0 * current_size) / @total_size
        percent_str = sprintf("%.1f", percent)
        eta_str = get_eta_str(current_size)
      else
        percent_str = "---.-"
        eta_str = "--:--"
      end

      speed_str = get_speed_str(current_size)
      current_size_str = format(current_size)
      total_size_str = format(@total_size)
      $stdout.printf("\rDownloading video data: %5s%% (%8s of %s) at %12s ETA %s ",
                     percent_str, current_size_str, total_size_str, speed_str, eta_str)
      $stdout.flush
    end

    private

    def get_eta_str(current_size)
      speed = calculate_speed(current_size)
      if speed
        rest_size = @total_size - current_size
        eta = (rest_size / speed).to_i
        eta_min, eta_sec = eta.divmod(60)
        if eta_min > 99
          "--:--"
        else
          sprintf "%02d:%02d", eta_min, eta_sec
        end
      else
        "--:--"
      end
    end

    def get_speed_str(current_size)
      speed = calculate_speed(current_size)
      "#{format(speed)}/sec"
    end

    def calculate_speed(current_size)
      elapsed = Time.now - @start_time
      if (current_size == 0) || (elapsed < Epsilon)
        nil
      else
        current_size.to_f / elapsed
      end
    end

    def format(value)
      if @unit == :byte
        format_byte(value)
      elsif @unit == :sec
        format_sec(value)
      else
        "#{value}"
      end
    end

    def format_byte(value)
      if value
        exp = (value > 0) ? Math.log(value, Size1K).to_i : 0
        suffix = "bKMGTPEZY"[exp]
        if exp == 0
          "#{value}#{suffix}"
        else
          converted = value.to_f / (Size1K ** exp)
          sprintf "%.2f%s", converted, suffix
        end
      else
        "N/A b"
      end
    end

    def format_sec(value)
      if value
        if value > 60
          converted = value.to_f / 60
          sprintf "%.1fmin", converted
        else
          sprintf "%.1fsec", value.to_f
        end
      else
        "N/A sec"
      end
    end
  end
end

今日はここまで！