class TorrentFile(object): def __init__(self, file_length, pieces_hash_array, piece_size, fs): self.logger = logging.getLogger(__name__) self.pieces_num = file_length / PIECE_SIZE self.piece_list = [] for piece in range(self.pieces_num): self.piece_list.append( Piece(piece, piece_size, fs, pieces_hash_array[piece])) #if have extra smaller pieces if file_length % PIECE_SIZE != 0: last_piece_size = file_length % piece_size self.piece_list.append( Piece(self.pieces_num, last_piece_size, fs, pieces_hash_array[self.pieces_num])) self.pieces_num += 1 self.bm = BitMap(self.pieces_num) self.logger.debug("init bm {}".format(self.bm)) def update_bitmap(self): for i in range(self.pieces_num): if self.piece_list[i].is_complete(): self.bm.set(i) else: self.bm.reset(i) def is_complete(self): # self.update_bitmap() return self.bm.all() ## return the missing pieces def missing_pieces(self): self.update_bitmap() missing_list = [] for i in range(self.pieces_num): if not self.bm.test(i): missing_list.append(self.piece_list[i]) return missing_list def set_piece(self, index, content): self.piece_list[index].set_block(content) self.piece_list[index].write_to_file() def get_info(self): for i in range(self.pieces_num): self.piece_list[i].get_info()
def test_thread(): global exitFlag exitFlag = 0 threads = [] threadID = 1 print "Starting Main Thread" # Fill the queue for i in range(3): queue = Queue.Queue(4) for j in range(4): queue.put(i*4 + j) pieces_queue_list.append(queue) for i in range(3): queue = Queue.Queue(4) data_queue_list.append(queue) for i , tName in enumerate(threadList): thread = myThread(threadID, tName, pieces_queue_list[i],data_queue_list[i]) thread.start() threads.append(thread) threadID += 1 # Wait for all threads to complete for t in threads: t.join() print "Exiting Main Thread" bm = BitMap(12) while 1: if bm.all(): print "receive all the pieces!!" break for q in data_queue_list: if not q.empty(): item = q.get() data = item['data'] index = item['index'] print "receive pieces:%s data:%s\n" % (data,data) bm.set(index)
class Piece(object): def __init__(self, index, piece_length, file_to_write, piece_hash): self.index = index self.piece_length = piece_length self.num_blocks = piece_length/BLOCK_SIZE self.file_to_write = file_to_write self.piece_hash = piece_hash self.block_list = [] for block in range(self.num_blocks): self.block_list.append(Block(block, BLOCK_SIZE)) if piece_length % BLOCK_SIZE != 0: last_block_length = piece_length % BLOCK_SIZE self.block_list.append(Block(self.num_blocks, last_block_length)) self.num_blocks = self.num_blocks + 1 self.bitmap = BitMap(self.num_blocks) # if self.index == 16: # print "init bm: ", self.bitmap def is_piece_full(self): for block in range(len(self.block_list)): if not self.block_list[block].is_block_full(): return False return True def fill_the_block(self, block_index, block_content): if self.block_list[block_index].is_block_full(): return self.block_list[block_index].write(block_content) self.bitmap.set(block_index) def write(self): if self.bitmap.all(): bytes_to_write = bytearray('') for block in self.block_list: bytes_to_write = bytes_to_write + block.block_content self.file_to_write.seek(self.index * self.piece_length) self.file_to_write.write(bytes_to_write)
class Piece(object): def __init__(self, piece_index, piece_size, fs, piece_hash): self.block_num = piece_size / BLOCK_SIZE self.block_list = [] self.piece_index = piece_index self.piece_size = piece_size self.piece_hash = piece_hash for block in range(self.block_num): b = Block(piece_index, block, BLOCK_SIZE) self.block_list.append(b) #if have extra smaller blocks if piece_size % BLOCK_SIZE != 0: last_block_size = piece_size % BLOCK_SIZE #print self.block_num,last_block_size self.block_list.append( Block(self.piece_index, self.block_num, last_block_size)) self.block_num += 1 self.bm = BitMap(self.block_num) self.file = fs return def __hash__(self): return hash((self.piece_index, self.piece_size, self.piece_hash)) def __eq__(self, other): return (self.piece_index, self.piece_size, self.piece_hash) == (other.piece_index, self.piece_size, self.piece_hash) def __ne__(self, other): return not (self == other) #check block's bitmap in disk and the whole SHA1 def is_complete(self): #check SHA1 self.file.seek(self.piece_index * PIECE_SIZE) content = self.file.read(self.piece_size) piecehash = hashlib.sha1(content).digest() # print piecehash, self.piece_hash # print self.piece_size return piecehash == self.piece_hash def written_piece_hash(self): self.file.seek(self.piece_index * PIECE_SIZE) content = self.file.read(self.piece_size) piecehash = hashlib.sha1(content).digest() return piecehash def update_bitmap(self): for i in range(len(self.block_list)): if self.block_list[i].is_complete(): self.bm.set(i) else: self.bm.reset(i) def get_info(self): print "piece index:", self.piece_index, " piece size:", self.piece_size #for i in range(self.block_num): # self.block_list[i].get_info() ## set every block def set_block(self, content): size = len(content) for block in self.block_list: #print "boffset",block.block_offset block.set_complete( content[block.block_offset:min(block.block_offset + BLOCK_SIZE, size)]) ## get the content from blocks and write to file def write_to_file(self): self.update_bitmap() if self.bm.all(): payload = bytearray('') for block in self.block_list: payload += block.payload self.file.seek(self.piece_index * PIECE_SIZE) self.file.write(payload) def clear_data(self): for block in self.block_list: block.payload = None def expire(self): for block in self.block_list: block.mark_missing() self.bm = BitMap(self.block_num)