#!/usr/bin/env ruby
module Bencode
  class 
<
< self

    def encode(object)
      case object
      when Hash
        encode_hash(object)
      when Integer
        encode_integer(object)
      when Array
        encode_array(object)
      when String
        encode_byte_string(object)
      end
    end

    def decode(data)
      case data[0].chr
        when ("0".."9")
          data,string = decode_byte_string(data)
          string
        when "l"
          data,array = decode_array(data[1..-1])
          array
        when "d"
          data,hash = decode_hash(data[1..-1])
          hash
        when "i"
          data,int = decode_integer(data)
          int
      end
    end

    private

    def decode_array(data)
      new_array = Array.new
      while data[0].chr != "e"
        case data[0].chr
        when ("0".."9")
          data,string = decode_byte_string(data)
          new_array.push(string)
        when "l"
          data,array = decode_array(data[1..-1])
          new_array.push(array)
        when "d"
          data,hash = decode_hash(data[1..-1])
          new_array.push(hash)
        when "i"
          data,int = decode_integer(data)
          new_array.push(int)
        end
      end
      return data[1..-1],new_array
    end

    def decode_hash(data)
      new_hash = Hash.new
      while data[0].chr != "e"
        data =~ /^(\d+):/
        key = data[$1.size+1...$1.size+1+$1.to_i]
        data = data[$1.size+1+$1.to_i..-1]
        value = case data[0].chr
        when ("0".."9")
          data,string = decode_byte_string(data)
          string
        when "l"
          data,array = decode_array(data[1..-1])
          array
        when "d"
          data,hash = decode_hash(data[1..-1])
          hash
        when "i"
          data,int = decode_integer(data)
          int
        end
        new_hash[key] = value
      end
      return data[1..-1],new_hash
    end

    def decode_byte_string(data)
        data =~ /(\d+):/
        bit_data = data[$1.size+1...$1.size+1+$1.to_i]
        data = data[$1.size+1+$1.to_i..-1]
        return data,bit_data
    end

    def decode_integer(data)
      data =~ /i(\d+)e/
      int = $1.to_i
      data.sub!('i'+$1+"e","")
      return data,int
    end

    def encode_byte_string(string)
      "#{string.size}:#{string}"
    end

    def encode_integer(int)
      "i#{int}e"
    end

    def encode_hash(hash)
      hash_string = "d"
      hash.each{|key,value|
        hash_string 
<
< encode_byte_string(key.to_s)
        hash_value = case value
        when Hash
          encode_hash(value)
        when Integer
          encode_integer(value)
        when Array
          encode_array(value)
        when String
          encode_byte_string(value)
        end
        hash_string 
<
< hash_value
      }
      hash_string+"e"
    end

    def encode_array(array)
      array_string = "l"
      array.each{|value|
        array_value = case value
        when Hash
          encode_hash(value)
        when Integer
          encode_integer(value)
        when Array
          encode_array(value)
        when String
          encode_byte_string(value)
        end
        array_string 
<
< array_value
      }
      array_string+"e"
    end
  end

end


if __FILE__ == $0

  require 'test/unit'
  class BencodeTest 
< Test::Unit::TestCase
    def setup
      @test_file = File.join(File.dirname(__FILE__),"test.torrent")
    end

    def test_decode_bencode_strings
      array = Bencode.decode("l4:spam4:eggsl4:spam4:eggsei67ee")
      assert_equal(["spam", "eggs", ["spam", "eggs"],67],array)
      hash = Bencode.decode("d3:cow3:moo4:spamd3:cow3:moo4:spami55eee")
      assert_equal({"cow"=>"moo", "spam"=>{"cow"=>"moo", "spam"=>55}},hash)
    end

    def test_encode_bencode_objects
      assert_equal("l4:spam4:eggsl4:spam4:eggsei67ee",Bencode.encode(["spam", "eggs", ["spam", "eggs"],67]))

      assert_equal("d3:cow3:moo4:spamd3:cow3:moo4:spami55eee",Bencode.encode({"cow"=>"moo", "spam"=>{"cow"=>"moo", "spam"=>55}}))
    end

    def test_decode_bencode_file
      torrent_data = File.read(@test_file)
      hash = Bencode.decode(torrent_data)
      assert_equal("UTF-8",hash["encoding"])
      assert(hash["info"]["name"])
      assert_equal("College Basketball - ''Colorado State at Utah'' (Recorded Jan 17, 2009, MTN).mpg",hash["info"]["name"])
      assert_equal(1232424648,hash["creation date"])
      assert_equal([["http://tracker.thepiratebay.org/announce"],
     ["udp://tracker.thepiratebay.org:80/announce"]],hash["announce-list"])
    end
  end
end