require 'socket'
require 'digest/sha1'

# decodes a .torrent file and torrent tracker responses
# format: an object
# 4 object types: list/dict/int/string
# list: 'l' + entries + 'e'
# dict: 'd' + key1 + val1 + key2 + val2 etc + 'e'
# int:  'i' + int.to_i(10) + 'e'
# string: length.to_i(10) + ':' + content
# list/dict keys/values are any object type
class Array
	def bencode
		'l' + map { |ee| ee.bencode }.join + 'e'
	end
end
class Hash
	def bencode
		# use @order if present, synched with current keys
		'd' + ((@order || []) & keys | keys).map { |k| k.bencode + self[k].bencode }.join + 'e'
	end
end
class Integer
	def bencode
		"i#{to_s}e"
	end
end
class String
	def bencode
		"#{length}:#{self}"
	end
	
	def bdecode(idx=0)
		ori_idx = idx
		case self[idx]
		when ?l
			ret = []
			idx += 1
			while self[idx] and self[idx] != ?e
				elem, idx = bdecode(idx)
				ret << elem
			end
			idx += 1
		when ?d
			ret = {}
			# the tracker need the hash of reencoding the torrent, and this needs to keep the order of fields
			class << ret ; attr_accessor :order end
			ret.order = []
			idx += 1
			while self[idx] and self[idx] != ?e
				key, idx = bdecode(idx)
				val, idx = bdecode(idx)
				ret.order << key
				ret[key] = val
			end
			idx += 1
		when ?i
			ret = ''
			idx += 1
			while self[idx] and self[idx] != ?e
				ret << self[idx]
				idx += 1
			end
			idx += 1
			ret = ret.to_i
		when ?0..?9
			len = ''
			while self[idx] and self[idx] != ?:
				len << self[idx]
				idx += 1
			end
			idx += 1
			len = len.to_i
			ret = self[idx, len]
			idx += len
		else raise "Invalid bstring character at index #{idx} #{inspect}"
		end
		ori_idx > 0 ? [ret, idx] : ret
	end
end

module Torrent
class Tracker
	attr_accessor :uploaded, :downloaded, :left
	attr_reader :info_hash, :id, :torrent

	# id must be 20 chars
	def initialize(torrent, id=nil, ip=nil)
		@torrent = torrent.bdecode
		raise "invalid torrent announce url" unless @torrent['announce'] =~ /http:\/\/(.*?)(?::(\d+))?(\/.*)/
		@host, @port, @path = $1, $2, $3
		@port = @port ? @port.to_i : 80
		@id = id || "-jjtorrent-#{rand(10000)}".ljust(20, '-')
		@info_hash = Digest::SHA1.digest(@torrent['info'].bencode)	# XXX need to keep the hash keys order from original
		@uploaded = @downloaded = 0
		@left = @torrent['info']['length'] || @torrent['info']['files'].inject(0) { |s, f| s + f['length'] }
		@ip = ip
	end

	def announce(event, port=nil)
		TCPSocket.open(@host, @port) { |s|
			# ip ||= s.addr[3]
			args = {}
			args['info_hash'] = @info_hash.unpack('C*').map { |c| '%%%02x' % c }.join
			args['peer_id']   = @id
			args['ip'] = @ip if @ip
			args['port'] = port if port
			args['uploaded'] = @uploaded
			args['downloaded'] = @downloaded
			args['left'] = @left
			args['event'] = event
			args['compact'] = 1
			#args['key'] = '%02x%02x%02x%02x' % [rand(256), rand(256), rand(256), rand(256)]

			req = @path + '?' + args.map { |k, v| "#{k}=#{v}" }.join('&')
			p req if $DEBUG
			s.write "GET #{req} HTTP/1.1\r\n" +
				"Host: #@host#{":#@port" if @port != 80}\r\n" +
				"Connection: close\r\n" +
				"\r\n"

			ans = s.gets
			case ans.split[1].to_i
			when 200
				puts "tracker ok" if $VERBOSE
			when 302
				puts "HTTP redirect"
				nil while ans = s.gets and ans !~ /^location.*http:\/\/(.*?)(?::(\d+))?(\/.*)/i
				@host, @port, @path = $1, $2, $3
				@port = @port ? @port.to_i : 80
				return announce(event, port)
			else raise "Invalid response #{ans} + #{s.read}"
			end
			nil while not s.gets.chomp.empty?	# ignore HTTP response headers
			s.read.bdecode				# decode body
		}
	end
end

class BitString
	attr_accessor :str
	def initialize(length)
		@str = 0.chr * 4 * ((length+31)/32)		# round to next 4byte boundary
	end

	def [](i)
		@str[i/8][7-(i%8)]
	end
	
	def []=(i, v)
		if v == 1
			@str[i/8] |= 1 << (7 - (i%8))
		else
			@str[i/8] &= 0xff - (1<<(7 - (i%8)))
		end
	end
end

class Peer
	TYPES = { :choke => 0, :unchoke => 1, :interested => 2, :notinterested => 3,
		:have => 4, :bitfield => 5, :request => 6, :piece => 7, :cancel => 8 }

	attr_reader :s, :id
	def initialize(ip, port, id)
		@ip, @port, @id = ip, port, id
		@s = nil
	end

	def handshake(info_hash, myid)
		# @s = TCPSocket.new(@ip, @port)
		success = false
		begin
		@s = Socket.new(Socket::AF_INET, Socket::SOCK_STREAM, 0)
		sockaddr = Socket.sockaddr_in(@port, @ip)
		begin
			@s.connect_nonblock(sockaddr)
		rescue Errno::EINPROGRESS
			raise "connect timeout" if not IO.select(nil, [@s], nil, 15)
			@s.connect_nonblock(sockaddr)
		end
# 3 getfl 4 setfs
# fugly, ain't it ?
@s.fcntl(4, (@s.fcntl(3) & (~IO::NONBLOCK)))

		# send protocol (1 byte length + string), extensions (8 bytes), hash (raw) and id (20 + 20 bytes)
		message = '' << 19 << 'BitTorrent protocol' << "\0\0\0\0\0\0\0\0" << info_hash << myid
		@s.write message
		len = @s.read(1)
		raise "Closed connexion" if not len
		raise "Invalid handshake response #{len.inspect}" if len[0] != 19
		proto = @s.read(19)
		raise "Invalid protocol #{proto.inspect}" if proto != 'BitTorrent protocol'
		ext = @s.read(8)	# ignored
		hash = @s.read(20)
		raise "Invalid peer hash #{hash.inspect}" if hash != info_hash	# TODO server
		id = @s.read(20)
		success = true
		ensure
			@s.close unless success
		end
	end

	def send(type, content='')
		type = TYPES[type] || type
		str = '' << type << content
		@s.write [str.length].pack('N') + str
	end

	def read
		raise 'peer ans timeout' if not IO.select([@s], nil, nil, 120)
		len = @s.read(4).unpack('N').first
		str = @s.read len
		[TYPES.index(str[0]), str[1..-1]]
	end
end

class Downloader
	def initialize(path)
		@tracker = Tracker.new(File.read(path), '-UT1600-42'.ljust(20, '-'))

		if not @files = @tracker.torrent['info']['files']
			@files = [{ 'path' => [@tracker.torrent['info']['name']], 'length' => @tracker.torrent['info']['length'] }]
		end

		@files.each { |f| f['path'] = f['path'].join('/').tr("\0/\n\r", '_') }
		@hashes = @tracker.torrent['info']['pieces'].scan(/.{20}/m)
		@piece_len = @tracker.torrent['info']['piece length']
		@full_len = @files.inject(0) { |len, f| len + f['length'] }
		@npieces = (@full_len + @piece_len - 1) / @piece_len
	end

	def check_files
		@my_bitfield = BitString.new(@npieces)
		off = foff = 0
		buf = ''
		good = bad = 0
		@files.each { |f|
			if not File.exist? f['path']
				puts "creating #{f['path'].inspect} (#{f['length']}o)" if $VERBOSE
				File.open(f['path'], 'wb') { |fd| fd.pos = f['length'] - 1 ; fd.write 0.chr }
			else
				puts "reading #{f['path'].inspect} (#{f['length']}o)" if $VERBOSE
				raise 'invalid size !' if File.size(f['path']) != f['length']
			end
			loop do
				tbuf = File.open(f['path'], 'rb') { |fd| fd.pos = foff ; fd.read(@piece_len - buf.length) }
				buf += tbuf
				off += tbuf.length
				foff += tbuf.length
				if buf.length == @piece_len
					h = Digest::SHA1.digest(buf)
					if h == @hashes[off/@piece_len - 1]
						puts "good piece at #{foff}" if $DEBUG
						@my_bitfield[off/@piece_len - 1] = 1
						good += 1
					else
						puts "bad piece at #{foff}" if $DEBUG
						bad += 1
					end
					buf = ''
				else
					foff = 0
					break
				end
			end
		}
		if not buf.empty?
			h = Digest::SHA1.digest(buf)
			if h == @hashes[off/@piece_len]
				puts "good piece at #{foff}" if $DEBUG
				@my_bitfield[off/@piece_len] = 1
				good += 1
			else
				puts "bad piece at #{foff}" if $DEBUG
				bad += 1
			end
		end
		puts @my_bitfield.str.unpack('C*').map{ |c| '%02x' % c }.join if $VERBOSE
		puts "check done : #{good} good, #{bad} bad pieces (got #{'%02f%%' % ((good * 100.0)/(good+bad))} of #{@npieces*@piece_len/1024/1024}Mo)"
	end

	def download_test(hpeers = [])
		if hpeers.empty?
			ans = @tracker.announce('started', 6999)
			return if not peers = ans['peers']
		else
			peers = hpeers.map { |s| ip, port = s.split ':' ; { 'ip' => ip, 'port' => port.to_i } }
		end

		if peers.kind_of? String
			peers = peers.scan(/....../).map { |e|
				ip, port = e.unpack('Nn')
				ip = [(ip >> 24), (ip >> 16), (ip >> 8), ip].map { |x| x & 255 } * '.'
				{ 'ip' => ip, 'port' => port }
			}
		end

		targetpiece = (@full_len-47) / @piece_len	# read the last 47 bytes ( = nfo for bsg)

		begin
			if not peer = peers.shift
				sleep 30 ; return download_test
			end

			puts "Connecting to peer #{peer['ip']}:#{peer['port']} #{peer['peer id'].inspect}"
			peer = Peer.new(peer['ip'], peer['port'], peer['peer id'])
			peer.handshake(@tracker.info_hash, @tracker.id)
			puts " Connected, waiting bitfield"
			type, body = peer.read
			bf = BitString.new(@npieces)
			bf.str = body.dup
			if bf[targetpiece] == 0
				raise ' has noes last piece'
			end
			puts "he has !"
			loop do
				peer.send(:interested)
				type, body = peer.read
				case type
				when :unchoke
					break
				when :have
					puts " he has #{body.unpack('N').first}"
				else
					raise " not unchoking: got #{type.inspect} #{body.inspect}"
				end
			end
			puts " unchoked ! requesting bytes"
			peer.send(:request, [targetpiece, ((@full_len - 47) % @piece_len), 47].pack('NNN'))
			type, body = peer.read
			if type != :piece
				raise "did not piece: got #{type.inspect} #{body.inspect}"
			else
				puts "got it ! #{type.inspect} piece offset #{body[0,8].unpack('NN').inspect} #{body[8..-1].inspect}"
			end
			peer.s.close
		rescue Object
			puts "#{$!.class}: #{$!}"
			sleep 1
			retry
		end
	end
end
end


if $0 == __FILE__
	Torrent::Downloader.new(ARGV.shift).download_test ARGV
end

__END__
4.1 accueil d'un nouveau téléchargeur
On est ici dans le cas ou un peer P qui n'a encore aucun morceau du fichier contacte un peer Q qui possède du contenu intéressant pour P et qui est disposé à l'offrir:

   1. Q transmet un message bitfield (type 5) indiquant quels tronçons il possède, <len=X+1><05><bi><tf><ie><ld><ta><il><le><=X>. Chaque bit du bitfield représente un tronçon (1° byte & 0x80 == disponibilité du 1° tronçon)
   2. P transmet un message interrested (type 2) pour signaler qu'il désire recevoir des données depuis Q. <00000001><02>
   3. Dès que Q est près à transmettre, il envoie le message unchoke <00000001><01>
   4. P peut maintenant demander des données. chaque message request doit indiquer le n° du tronçon concerné (index dans le bitfield), et un intervalle de bytes de ce tronçon à transmettre (offset et longueur). Le protocole recommande des requêtes de 32K bytes. Les messages à envoyer successivement pour récupérer l'entièreté du 3° tronçon de 256K seraient donc:
          * --> request <0000000d><06><00000003><00000000><00008000>
          * <-- piece    <00008009><07><00000003><00000000> 32K de données
          * --> request <0000000d><06><00000003><00008000><00008000>
          * <-- piece    <00008009><07><00000003><00008000> 32K de données
          * --> request <0000000d><06><00000003><00010000><00008000>
          * ...
          * <-- piece    <00008009><07><00000003><00038000> 32K de données

4.2 un tronçon a été complètement reçu
Nous sommes ici dans le cas où un peer P vient de recevoir le dernier bloc complétant le tronçon #i alors qu'il est connecté aux peers Q1..Qn.

   1. P va calculer le hash SHA1 du tronçon reçu et s'assurer qu'il correspond bien au hash stocké dans le fichier info, c.à.d. aux bytes [i*20..19+i*20] de la chaîne info::pieces
   2. Si les données concordent avec le hashage annoncé, P peut ajouter les données au fichier récupéré de manière permanente.
   3. P annonce à ses voisins qu'il possède maintenant le tronçon numéro i à l'aide du message have. <00000005><04><iiiiiiii>
   4. Les voisins de P vont, suite à cette annonce, signaler leur (dés)intérêt par un des messages interested <00000001><02> ou not interested <00000001><03>.
   5. Si P est disposé à accepter de nouvelles connexions en upload, il transmettra un message unchoke <00000001><01> en réponse à interested. Le message choke <00000001><00> lui permet au contraire de mettre en attente un peer.