from io import BytesIO from collections import defaultdict class LecteurDeBits: def __init__(self, entree : bytes): self.octets = BytesIO(entree) self.bits_non_lus = 0 self.taille_bits_non_lus = 0 def lire_bits(self, nombre_bits : int) -> int: while self.taille_bits_non_lus < nombre_bits: prochain_octet = self.octets.read(1) if not prochain_octet: raise EOFError self.bits_non_lus |= prochain_octet[0] << self.taille_bits_non_lus self.taille_bits_non_lus += 8 masque_bits_lus = (1 << nombre_bits) - 1 bits_lus = self.bits_non_lus & masque_bits_lus self.bits_non_lus >>= nombre_bits self.taille_bits_non_lus -= nombre_bits return bits_lus def lire_octets(self, nombre_octets : int) -> bytes: self.aligner_bits_sur_octet() octets_lus = self.octets.read(nombre_octets) if len(octets_lus) < nombre_octets: raise EOFError return octets_lus def aligner_bits_sur_octet(self): self.bits_non_lus = 0 self.taille_bits_non_lus = 0 class LecteurDeBitsRangeCode: def __init__(self, entree : bytes): self.octets = BytesIO(entree) self.taille_code = 0xffffffff self.code = 0 assert self.octets.read(1)[0] == 0 for position in range(4): self.code = (self.code << 8) | self.octets.read(1)[0] assert self.code < self.taille_code def reprendre_des_bits_si_besoin(self): if self.taille_code <= 0xffffff: self.taille_code <<= 8 self.code <<= 8 self.code |= self.octets.read(1)[0] class RangeDecoder: def __init__(self, lecteur_de_bits : LecteurDeBitsRangeCode, probabilite_initiale : int = 0x400): self.lecteur_de_bits : LecteurDeBitsRangeCode = lecteur_de_bits self.probabilite : int = probabilite_initiale # 0 - 0x800 def lire_bit(self, utiliser_probas = True): if utiliser_probas: milieu_du_code_pondere = self.lecteur_de_bits.taille_code // 0x800 * self.probabilite else: milieu_du_code_pondere = self.lecteur_de_bits.taille_code // 2 if self.lecteur_de_bits.code < milieu_du_code_pondere: bit_lu = 0 self.probabilite += (0x800 - self.probabilite) // 32 self.lecteur_de_bits.taille_code = milieu_du_code_pondere else: bit_lu = 1 self.probabilite -= self.probabilite // 32 self.lecteur_de_bits.code -= milieu_du_code_pondere if utiliser_probas: self.lecteur_de_bits.taille_code -= milieu_du_code_pondere else: self.lecteur_de_bits.taille_code = milieu_du_code_pondere self.lecteur_de_bits.reprendre_des_bits_si_besoin() return bit_lu class DecodeurLZMASPD: def decode(self, entree : bytes) -> bytes: self.flux_decompresse = b'' self.dernieres_distances : List[int] = [0] * 4 lecteur_de_bits = LecteurDeBits(entree) properties = lecteur_de_bits.lire_bits(8) self.literal_context_bits = properties % 9 literal_position_bits = (properties // 9) % 5 position_bits = properties // 9 // 5 if position_bits > 4: raise ValueError('LZMA invalid') self.taille_fenetre = lecteur_de_bits.lire_bits(32) self.uncompressed_size = lecteur_de_bits.lire_bits(64) lecteur_de_bits = LecteurDeBitsRangeCode(entree[13:]) self.nom_vers_range_decodeur : Dict[tuple, RangeDecoder] = defaultdict(lambda: RangeDecoder(lecteur_de_bits)) self.state = 0 self.masque_pos_state = (1 << position_bits) - 1 self.masque_lit_state = (1 << literal_position_bits) - 1 while len(self.flux_decompresse) < self.uncompressed_size: try: pos_state = len(self.flux_decompresse) & self.masque_pos_state bit_choice = self.nom_vers_range_decodeur[('IsMatch', self.state, pos_state)].lire_bit() if bit_choice == 0: self.LITERAL() elif bit_choice == 1: is_rep = self.nom_vers_range_decodeur[('IsRep', self.state)].lire_bit() if is_rep == 0: self.MATCH() elif is_rep == 1: is_rep0 = self.nom_vers_range_decodeur[('IsRepG0', self.state)].lire_bit() if is_rep0 == 0: is_rep0_long = self.nom_vers_range_decodeur[('IsRep0Long', self.state, pos_state)].lire_bit() if is_rep0_long == 0: self.SHORTREP() elif is_rep0_long == 1: self.LONGREP(0) elif is_rep0 == 1: is_rep1 = self.nom_vers_range_decodeur[('IsRepG1', self.state)].lire_bit() if is_rep1 == 0: self.LONGREP(1) elif is_rep1 == 1: is_rep2 = self.nom_vers_range_decodeur[('IsRepG2', self.state)].lire_bit() if is_rep2 == 0: self.LONGREP(2) elif is_rep2 == 1: self.LONGREP(3) except EOFError: break return self.flux_decompresse def LITERAL(self): dernier_octet_decompresse = self.flux_decompresse[-1] if self.flux_decompresse else 0 octet_lu = self.bit_tree_decode(('LiteralNormal', len(self.flux_decompresse) & self.masque_lit_state, # total_pos dernier_octet_decompresse >> (8 - self.literal_context_bits), # prev_byte ), None, 8, use_pos_state = False) self.flux_decompresse += bytes([octet_lu]) if self.state > 9: self.state -= 6 elif self.state > 3: self.state -= 3 else: self.state = 0 def MATCH(self): match_len = 2 + self.len_decode('LenDecoder') pos_slot = self.bit_tree_decode(('PosSlot', min(5, match_len)), None, 6, use_pos_state = False) if pos_slot >= 4: num_direct_bits = (pos_slot >> 1) - 1 distance = (2 | (pos_slot & 1)) << num_direct_bits if pos_slot < 14: distance += self.bit_tree_decode('SpecPos', None, num_direct_bits + (distance - pos_slot - 1), use_pos_state = False, reverse = True, debut_bit_tree = distance - pos_slot - 1) else: distance += self.bit_tree_decode('AlignFixed', None, num_direct_bits - 4, utiliser_probas = False) << 4 distance += self.bit_tree_decode('Align', None, 4, use_pos_state = False, reverse = True) else: distance = pos_slot self.dernieres_distances.append(distance) if distance == 0xffffffff: raise EOFError assert distance < len(self.flux_decompresse) assert distance < self.taille_fenetre self.repeter_donnees(distance, match_len) if self.state < 7: self.state = 7 else: self.state = 10 def SHORTREP(self): # Réutiliser la dernière distance pour un octet if self.state < 7: self.state = 9 else: self.state = 11 self.repeter_donnees(self.dernieres_distances[-1], 1) def LONGREP(self, num): # Réutiliser l'une des dernières distances pour une taille donnée match_len = 2 + self.len_decode('RepLenDecoder') self.dernieres_distances.append(self.dernieres_distances.pop(-(1 + num))) distance = self.dernieres_distances[-1] self.repeter_donnees(distance, match_len) if self.state < 7: self.state = 8 else: self.state = 11 def repeter_donnees(self, distance, match_len): debut_slice = len(self.flux_decompresse) - (1 + distance) fin_slice = match_len a_repeter = self.flux_decompresse[debut_slice:debut_slice + fin_slice] fin_slice -= len(self.flux_decompresse) - debut_slice while fin_slice > 0: a_repeter += self.flux_decompresse[debut_slice:debut_slice + fin_slice] fin_slice -= min(len(self.flux_decompresse), debut_slice + fin_slice) - debut_slice self.flux_decompresse += a_repeter def len_decode(self, len_decoder_name): if self.nom_vers_range_decodeur[('LenChoice', len_decoder_name)].lire_bit() == 0: return self.bit_tree_decode('LenLow', len_decoder_name, 3) else: if self.nom_vers_range_decodeur[('LenChoice2', len_decoder_name)].lire_bit() == 0: return (1 << 3) + self.bit_tree_decode('LenMid', len_decoder_name, 3) else: return (1 << 4) + self.bit_tree_decode('LenHigh', len_decoder_name, 8, use_pos_state = False) def bit_tree_decode(self, bit_tree_decoder_name, len_decoder_name, num_bits, use_pos_state = True, utiliser_probas = True, reverse = False, bit_tree_lu = 0, debut_bit_tree = 0): for position_bit in range(debut_bit_tree, num_bits): bit_lu = self.nom_vers_range_decodeur[( bit_tree_decoder_name, len_decoder_name, bit_tree_lu, position_bit, (len(self.flux_decompresse) & self.masque_pos_state) if use_pos_state else None, )].lire_bit(utiliser_probas = utiliser_probas) if not reverse: bit_tree_lu <<= 1 bit_tree_lu |= bit_lu else: bit_tree_lu |= bit_lu << (position_bit - debut_bit_tree) return bit_tree_lu