<s> from __future__ import division , unicode_literals <EOL> import os <EOL> import hashlib <EOL> import logging <EOL> from collections import defaultdict <EOL> from . bencode import bencode , bdecode <EOL> from . humanize import humanize_bytes <EOL> from . utils import is_unsplitable , get_root_of_unsplitable , Pieces <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> class Color : <EOL> BLACK = '<STR_LIT>' <EOL> RED = '<STR_LIT>' <EOL> GREEN = '<STR_LIT>' <EOL> YELLOW = '<STR_LIT>' <EOL> BLUE = '<STR_LIT>' <EOL> PINK = '<STR_LIT>' <EOL> CYAN = '<STR_LIT>' <EOL> WHITE = '<STR_LIT>' <EOL> ENDC = '<STR_LIT>' <EOL> COLOR_OK = Color . GREEN <EOL> COLOR_MISSING_FILES = Color . RED <EOL> COLOR_ALREADY_SEEDING = Color . BLUE <EOL> COLOR_FOLDER_EXIST_NOT_SEEDING = Color . YELLOW <EOL> COLOR_FAILED_TO_ADD_TO_CLIENT = Color . PINK <EOL> class Status : <EOL> OK = <NUM_LIT:0> <EOL> MISSING_FILES = <NUM_LIT:1> <EOL> ALREADY_SEEDING = <NUM_LIT:2> <EOL> FOLDER_EXIST_NOT_SEEDING = <NUM_LIT:3> <EOL> FAILED_TO_ADD_TO_CLIENT = <NUM_LIT:4> <EOL> status_messages = { <EOL> Status . OK : '<STR_LIT>' % ( COLOR_OK , Color . ENDC ) , <EOL> Status . MISSING_FILES : '<STR_LIT>' % ( COLOR_MISSING_FILES , Color . ENDC ) , <EOL> Status . ALREADY_SEEDING : '<STR_LIT>' % ( COLOR_ALREADY_SEEDING , Color . ENDC ) , <EOL> Status . FOLDER_EXIST_NOT_SEEDING : '<STR_LIT>' % ( COLOR_FOLDER_EXIST_NOT_SEEDING , Color . ENDC ) , <EOL> Status . FAILED_TO_ADD_TO_CLIENT : '<STR_LIT>' % ( COLOR_FAILED_TO_ADD_TO_CLIENT , Color . ENDC ) , <EOL> } <EOL> CHUNK_SIZE = <NUM_LIT> <EOL> class UnknownLinkTypeException ( Exception ) : <EOL> pass <EOL> class IllegalPathException ( Exception ) : <EOL> pass <EOL> class AutoTorrent ( object ) : <EOL> def __init__ ( self , db , client , store_path , add_limit_size , add_limit_percent , delete_torrents , link_type = '<STR_LIT>' ) : <EOL> self . db = db <EOL> self . client = client <EOL> self . store_path = store_path <EOL> self . add_limit_size = add_limit_size <EOL> self . add_limit_percent = add_limit_percent <EOL> self . delete_torrents = delete_torrents <EOL> self . link_type = link_type <EOL> self . torrents_seeded = set ( ) <EOL> def try_decode ( self , value ) : <EOL> try : <EOL> return value . decode ( '<STR_LIT:utf-8>' ) <EOL> except UnicodeDecodeError : <EOL> logger . debug ( '<STR_LIT>' % value ) <EOL> return value . decode ( '<STR_LIT>' ) <EOL> def is_legal_path ( self , path ) : <EOL> for p in path : <EOL> if p in [ '<STR_LIT:.>' , '<STR_LIT:..>' ] or '<STR_LIT:/>' in p : <EOL> return False <EOL> return True <EOL> def populate_torrents_seeded ( self ) : <EOL> """<STR_LIT>""" <EOL> self . torrents_seeded = set ( x . lower ( ) for x in self . client . get_torrents ( ) ) <EOL> def get_info_hash ( self , torrent ) : <EOL> """<STR_LIT>""" <EOL> return hashlib . sha1 ( bencode ( torrent [ b'<STR_LIT:info>' ] ) ) . hexdigest ( ) <EOL> def find_hash_checks ( self , torrent , result ) : <EOL> """<STR_LIT>""" <EOL> modified_result = False <EOL> pieces = Pieces ( torrent ) <EOL> if self . db . hash_slow_mode : <EOL> logger . info ( '<STR_LIT>' ) <EOL> self . db . build_hash_size_table ( ) <EOL> start_size = <NUM_LIT:0> <EOL> end_size = <NUM_LIT:0> <EOL> logger . info ( '<STR_LIT>' ) <EOL> for f in result : <EOL> start_size = end_size <EOL> end_size += f [ '<STR_LIT>' ] <EOL> if f [ '<STR_LIT>' ] : <EOL> continue <EOL> files_to_check = [ ] <EOL> logger . debug ( '<STR_LIT>' ) <EOL> if self . db . hash_size_mode : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> files_to_check += self . db . find_hash_size ( f [ '<STR_LIT>' ] ) <EOL> if self . db . hash_name_mode : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> name = f [ '<STR_LIT:path>' ] [ - <NUM_LIT:1> ] <EOL> files_to_check += self . db . find_hash_name ( name ) <EOL> if self . db . hash_slow_mode : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> files_to_check += self . db . find_hash_varying_size ( f [ '<STR_LIT>' ] ) <EOL> logger . debug ( '<STR_LIT>' % len ( files_to_check ) ) <EOL> checked_files = set ( ) <EOL> for db_file in files_to_check : <EOL> if db_file in checked_files : <EOL> logger . debug ( '<STR_LIT>' % db_file ) <EOL> checked_files . add ( db_file ) <EOL> logger . info ( '<STR_LIT>' % db_file ) <EOL> match_start , match_end = pieces . match_file ( db_file , start_size , end_size ) <EOL> logger . info ( '<STR_LIT>' % ( db_file , match_start , match_end ) ) <EOL> if match_start or match_end : <EOL> size = os . path . getsize ( db_file ) <EOL> if size != f [ '<STR_LIT>' ] : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> if match_start and match_end : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> modification_point = pieces . find_piece_breakpoint ( db_file , start_size , end_size ) <EOL> elif match_start : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> modification_point = min ( f [ '<STR_LIT>' ] , size ) <EOL> elif match_end : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> modification_point = <NUM_LIT:0> <EOL> if size > f [ '<STR_LIT>' ] : <EOL> modification_action = '<STR_LIT>' <EOL> else : <EOL> modification_action = '<STR_LIT>' <EOL> f [ '<STR_LIT>' ] = False <EOL> f [ '<STR_LIT>' ] = ( '<STR_LIT>' , modification_action , modification_point ) <EOL> modified_result = True <EOL> else : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> f [ '<STR_LIT>' ] = True <EOL> f [ '<STR_LIT>' ] = db_file <EOL> break <EOL> return modified_result , result <EOL> def index_torrent ( self , torrent ) : <EOL> """<STR_LIT>""" <EOL> torrent_name = torrent [ b'<STR_LIT:info>' ] [ b'<STR_LIT:name>' ] <EOL> logger . debug ( '<STR_LIT>' % ( torrent_name , ) ) <EOL> torrent_name = self . try_decode ( torrent_name ) <EOL> if not self . is_legal_path ( [ torrent_name ] ) : <EOL> raise IllegalPathException ( '<STR_LIT>' % torrent_name ) <EOL> logger . info ( '<STR_LIT>' % torrent_name ) <EOL> if self . db . exact_mode : <EOL> prefix = '<STR_LIT:d>' if b'<STR_LIT>' in torrent [ b'<STR_LIT:info>' ] else '<STR_LIT:f>' <EOL> paths = self . db . find_exact_file_path ( prefix , torrent_name ) <EOL> if paths : <EOL> for path in paths : <EOL> logger . debug ( '<STR_LIT>' % path ) <EOL> if prefix == '<STR_LIT:f>' : <EOL> logger . info ( '<STR_LIT>' ) <EOL> size = os . path . getsize ( path ) <EOL> if torrent [ b'<STR_LIT:info>' ] [ b'<STR_LIT>' ] != size : <EOL> continue <EOL> return { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : os . path . dirname ( path ) , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT>' : path , <EOL> '<STR_LIT>' : size , <EOL> '<STR_LIT:path>' : [ torrent_name ] , <EOL> '<STR_LIT>' : True , <EOL> } ] } <EOL> else : <EOL> result = [ ] <EOL> for f in torrent [ b'<STR_LIT:info>' ] [ b'<STR_LIT>' ] : <EOL> orig_path = [ self . try_decode ( x ) for x in f [ b'<STR_LIT:path>' ] ] <EOL> p = os . path . join ( path , * orig_path ) <EOL> if not os . path . isfile ( p ) : <EOL> logger . debug ( '<STR_LIT>' % p ) <EOL> break <EOL> size = os . path . getsize ( p ) <EOL> if size != f [ b'<STR_LIT>' ] : <EOL> logger . debug ( '<STR_LIT>' % ( p , size , f [ b'<STR_LIT>' ] ) ) <EOL> break <EOL> result . append ( { <EOL> '<STR_LIT>' : p , <EOL> '<STR_LIT>' : f [ b'<STR_LIT>' ] , <EOL> '<STR_LIT:path>' : orig_path , <EOL> '<STR_LIT>' : True , <EOL> } ) <EOL> else : <EOL> logger . info ( '<STR_LIT>' ) <EOL> return { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : path , <EOL> '<STR_LIT>' : result } <EOL> result = [ ] <EOL> if b'<STR_LIT>' in torrent [ b'<STR_LIT:info>' ] : <EOL> files_sorted = { } <EOL> files = { } <EOL> if b'<STR_LIT>' in torrent [ b'<STR_LIT:info>' ] : <EOL> i = <NUM_LIT:0> <EOL> path_files = defaultdict ( list ) <EOL> for f in torrent [ b'<STR_LIT:info>' ] [ b'<STR_LIT>' ] : <EOL> logger . debug ( '<STR_LIT>' % ( f , ) ) <EOL> orig_path = [ self . try_decode ( x ) for x in f [ b'<STR_LIT:path>' ] if x ] <EOL> if not self . is_legal_path ( orig_path ) : <EOL> raise IllegalPathException ( '<STR_LIT>' % orig_path ) <EOL> path = [ torrent_name ] + orig_path <EOL> name = path . pop ( ) <EOL> path_files [ os . path . join ( * path ) ] . append ( { <EOL> '<STR_LIT:path>' : orig_path , <EOL> '<STR_LIT>' : f [ b'<STR_LIT>' ] , <EOL> } ) <EOL> files_sorted [ '<STR_LIT:/>' . join ( orig_path ) ] = i <EOL> i += <NUM_LIT:1> <EOL> if self . db . unsplitable_mode : <EOL> unsplitable_paths = set ( ) <EOL> for path , files in path_files . items ( ) : <EOL> if is_unsplitable ( f [ '<STR_LIT:path>' ] [ - <NUM_LIT:1> ] for f in files ) : <EOL> path = path . split ( os . sep ) <EOL> name = get_root_of_unsplitable ( path ) <EOL> if not name : <EOL> continue <EOL> while path [ - <NUM_LIT:1> ] != name : <EOL> path . pop ( ) <EOL> unsplitable_paths . add ( os . path . join ( * path ) ) <EOL> for path , files in path_files . items ( ) : <EOL> if self . db . unsplitable_mode : <EOL> path = path . split ( os . sep ) <EOL> while path and os . path . join ( * path ) not in unsplitable_paths : <EOL> path . pop ( ) <EOL> else : <EOL> path = None <EOL> if path : <EOL> name = path [ - <NUM_LIT:1> ] <EOL> for f in files : <EOL> actual_path = self . db . find_unsplitable_file_path ( name , f [ '<STR_LIT:path>' ] , f [ '<STR_LIT>' ] ) <EOL> f [ '<STR_LIT>' ] = actual_path <EOL> f [ '<STR_LIT>' ] = actual_path is not None <EOL> result += files <EOL> else : <EOL> for f in files : <EOL> actual_path = self . db . find_file_path ( f [ '<STR_LIT:path>' ] [ - <NUM_LIT:1> ] , f [ '<STR_LIT>' ] ) <EOL> f [ '<STR_LIT>' ] = actual_path <EOL> f [ '<STR_LIT>' ] = actual_path is not None <EOL> result += files <EOL> result = sorted ( result , key = lambda x : files_sorted [ '<STR_LIT:/>' . join ( x [ '<STR_LIT:path>' ] ) ] ) <EOL> else : <EOL> length = torrent [ b'<STR_LIT:info>' ] [ b'<STR_LIT>' ] <EOL> actual_path = self . db . find_file_path ( torrent_name , length ) <EOL> result . append ( { <EOL> '<STR_LIT>' : actual_path , <EOL> '<STR_LIT>' : length , <EOL> '<STR_LIT:path>' : [ torrent_name ] , <EOL> '<STR_LIT>' : actual_path is not None , <EOL> } ) <EOL> mode = '<STR_LIT>' <EOL> if self . db . hash_mode : <EOL> modified_result , result = self . find_hash_checks ( torrent , result ) <EOL> if modified_result : <EOL> mode = '<STR_LIT>' <EOL> return { '<STR_LIT>' : mode , '<STR_LIT>' : result } <EOL> def parse_torrent ( self , torrent ) : <EOL> """<STR_LIT>""" <EOL> files = self . index_torrent ( torrent ) <EOL> found_size , missing_size = <NUM_LIT:0> , <NUM_LIT:0> <EOL> for f in files [ '<STR_LIT>' ] : <EOL> if f [ '<STR_LIT>' ] or f . get ( '<STR_LIT>' ) : <EOL> found_size += f [ '<STR_LIT>' ] <EOL> else : <EOL> missing_size += f [ '<STR_LIT>' ] <EOL> return found_size , missing_size , files <EOL> def link_files ( self , destination_path , files ) : <EOL> """<STR_LIT>""" <EOL> if not os . path . isdir ( destination_path ) : <EOL> os . makedirs ( destination_path ) <EOL> for f in files : <EOL> if f [ '<STR_LIT>' ] : <EOL> destination = os . path . join ( destination_path , * f [ '<STR_LIT:path>' ] ) <EOL> file_path = os . path . dirname ( destination ) <EOL> if not os . path . isdir ( file_path ) : <EOL> logger . debug ( '<STR_LIT>' % file_path ) <EOL> os . makedirs ( file_path ) <EOL> logger . debug ( '<STR_LIT>' % ( self . link_type , f [ '<STR_LIT>' ] , destination ) ) <EOL> if self . link_type == '<STR_LIT>' : <EOL> os . symlink ( f [ '<STR_LIT>' ] , destination ) <EOL> elif self . link_type == '<STR_LIT>' : <EOL> os . link ( f [ '<STR_LIT>' ] , destination ) <EOL> else : <EOL> raise UnknownLinkTypeException ( '<STR_LIT>' % self . link_type ) <EOL> def rewrite_hashed_files ( self , destination_path , files ) : <EOL> """<STR_LIT>""" <EOL> if not os . path . isdir ( destination_path ) : <EOL> os . makedirs ( destination_path ) <EOL> for f in files : <EOL> if not f [ '<STR_LIT>' ] and '<STR_LIT>' in f : <EOL> destination = os . path . join ( destination_path , * f [ '<STR_LIT:path>' ] ) <EOL> file_path = os . path . dirname ( destination ) <EOL> if not os . path . isdir ( file_path ) : <EOL> logger . debug ( '<STR_LIT>' % file_path ) <EOL> os . makedirs ( file_path ) <EOL> logger . debug ( '<STR_LIT>' % ( f [ '<STR_LIT>' ] , destination ) ) <EOL> _ , modification_action , modification_point = f [ '<STR_LIT>' ] <EOL> current_size = os . path . getsize ( f [ '<STR_LIT>' ] ) <EOL> expected_size = f [ '<STR_LIT>' ] <EOL> diff = abs ( current_size - expected_size ) <EOL> modified = False <EOL> bytes_written = <NUM_LIT:0> <EOL> with open ( destination , '<STR_LIT:wb>' ) as output_fp : <EOL> with open ( f [ '<STR_LIT>' ] , '<STR_LIT:rb>' ) as input_fp : <EOL> logger . debug ( '<STR_LIT>' % ( f [ '<STR_LIT>' ] , destination , modification_point ) ) <EOL> while True : <EOL> if not modified and bytes_written == modification_point : <EOL> logger . debug ( '<STR_LIT>' % ( modification_action , diff ) ) <EOL> modified = True <EOL> if modification_action == '<STR_LIT>' : <EOL> seek_point = bytes_written + diff <EOL> logger . debug ( '<STR_LIT>' % ( seek_point , ) ) <EOL> input_fp . seek ( seek_point ) <EOL> elif modification_action == '<STR_LIT>' : <EOL> logger . debug ( '<STR_LIT>' % diff ) <EOL> while diff > <NUM_LIT:0> : <EOL> write_bytes = min ( CHUNK_SIZE , diff ) <EOL> output_fp . write ( b'<STR_LIT:\x00>' * write_bytes ) <EOL> diff -= write_bytes <EOL> read_bytes = CHUNK_SIZE <EOL> if not modified : <EOL> read_bytes = min ( read_bytes , modification_point - bytes_written ) <EOL> logger . debug ( '<STR_LIT>' % ( read_bytes , ) ) <EOL> data = input_fp . read ( read_bytes ) <EOL> if not data : <EOL> break <EOL> output_fp . write ( data ) <EOL> bytes_written += read_bytes <EOL> logger . debug ( '<STR_LIT>' ) <EOL> def handle_torrentfile ( self , path , dry_run = False ) : <EOL> """<STR_LIT>""" <EOL> logger . info ( '<STR_LIT>' % path ) <EOL> torrent = self . open_torrentfile ( path ) <EOL> if self . check_torrent_in_client ( torrent ) : <EOL> self . print_status ( Status . ALREADY_SEEDING , path , '<STR_LIT>' ) <EOL> if self . delete_torrents : <EOL> logger . info ( '<STR_LIT>' % path ) <EOL> os . remove ( path ) <EOL> return Status . ALREADY_SEEDING <EOL> found_size , missing_size , files = self . parse_torrent ( torrent ) <EOL> missing_percent = ( missing_size / ( found_size + missing_size ) ) * <NUM_LIT:100> <EOL> found_percent = <NUM_LIT:100> - missing_percent <EOL> would_not_add = missing_size and missing_percent > self . add_limit_percent or missing_size > self . add_limit_size <EOL> if dry_run : <EOL> return found_size , missing_size , would_not_add , [ f [ '<STR_LIT>' ] for f in files [ '<STR_LIT>' ] if f . get ( '<STR_LIT>' ) ] <EOL> if would_not_add : <EOL> logger . info ( '<STR_LIT>' % ( path , found_percent , humanize_bytes ( missing_size ) ) ) <EOL> self . print_status ( Status . MISSING_FILES , path , '<STR_LIT>' % ( found_percent , humanize_bytes ( missing_size ) ) ) <EOL> return Status . MISSING_FILES <EOL> if files [ '<STR_LIT>' ] == '<STR_LIT>' or files [ '<STR_LIT>' ] == '<STR_LIT>' : <EOL> logger . info ( '<STR_LIT>' ) <EOL> destination_path = os . path . join ( self . store_path , os . path . splitext ( os . path . basename ( path ) ) [ <NUM_LIT:0> ] ) <EOL> if os . path . isdir ( destination_path ) : <EOL> logger . info ( '<STR_LIT>' % destination_path ) <EOL> self . print_status ( Status . FOLDER_EXIST_NOT_SEEDING , path , '<STR_LIT>' ) <EOL> return Status . FOLDER_EXIST_NOT_SEEDING <EOL> self . link_files ( destination_path , files [ '<STR_LIT>' ] ) <EOL> elif files [ '<STR_LIT>' ] == '<STR_LIT>' : <EOL> logger . info ( '<STR_LIT>' ) <EOL> destination_path = files [ '<STR_LIT>' ] <EOL> fast_resume = True <EOL> if files [ '<STR_LIT>' ] == '<STR_LIT>' : <EOL> fast_resume = False <EOL> logger . info ( '<STR_LIT>' ) <EOL> self . rewrite_hashed_files ( destination_path , files [ '<STR_LIT>' ] ) <EOL> if self . delete_torrents : <EOL> logger . info ( '<STR_LIT>' % path ) <EOL> os . remove ( path ) <EOL> if self . client . add_torrent ( torrent , destination_path , files [ '<STR_LIT>' ] , fast_resume ) : <EOL> self . print_status ( Status . OK , path , '<STR_LIT>' ) <EOL> return Status . OK <EOL> else : <EOL> self . print_status ( Status . FAILED_TO_ADD_TO_CLIENT , path , '<STR_LIT>' ) <EOL> return Status . FAILED_TO_ADD_TO_CLIENT <EOL> def check_torrent_in_client ( self , torrent ) : <EOL> """<STR_LIT>""" <EOL> info_hash = self . get_info_hash ( torrent ) <EOL> return info_hash in self . torrents_seeded <EOL> def open_torrentfile ( self , path ) : <EOL> """<STR_LIT>""" <EOL> with open ( path , '<STR_LIT:rb>' ) as f : <EOL> return bdecode ( f . read ( ) ) <EOL> def print_status ( self , status , torrentfile , message ) : <EOL> print ( '<STR_LIT>' % ( '<STR_LIT>' % status_messages [ status ] , os . path . splitext ( os . path . basename ( torrentfile ) ) [ <NUM_LIT:0> ] , message ) ) </s>
