| 170 | |
| 171 | |
| 172 | class PycHandler(BaseHandler): |
| 173 | def __init__(self, file_path): |
| 174 | super().__init__(file_path) |
| 175 | with open(file_path, 'rb') as f: |
| 176 | self.data = f.read() |
| 177 | |
| 178 | def verify(self): |
| 179 | # print(repr(self.data[:10])) # TODO: Make this work with all pyc versions |
| 180 | return self.data[:4] == b'o\r\r\n' |
| 181 | |
| 182 | def handle(self): |
| 183 | '''if not self.verify(): |
| 184 | pass |
| 185 | raise ValueError('Invalid pyc file')''' |
| 186 | |
| 187 | # Disassemble the pyc file with pycdas executable in config.PYCDAS_PATH |
| 188 | pycdas_path = config.PYCDAS_PATH |
| 189 | if not os.path.exists(pycdas_path): |
| 190 | raise FileNotFoundError('pycdas executable not found') |
| 191 | logger.info(f'Converting pyc file to bytecode : {self.file_path}') |
| 192 | bytecode_text = subprocess.check_output( |
| 193 | [pycdas_path, self.file_path], stderr=subprocess.DEVNULL) |
| 194 | |
| 195 | # Split disassembled python code into functions using the PycCodeObject class |
| 196 | bytecode_text = bytecode_text.decode().split('\n') |
| 197 | file_as_code_object = PycCodeObject(bytecode_text[1:]) |
| 198 | |
| 199 | # For optimization, now we need to try to disassemble the pyc back to source code using pycdc, and then only fix the functions that are not correctly disassembled |
| 200 | # This is because pycdas does not always disassemble the code correctly |
| 201 | pycdc_path = config.PYCDC_PATH |
| 202 | if not os.path.exists(pycdc_path): |
| 203 | raise FileNotFoundError('pycdc executable not found') |
| 204 | logger.info(f'Converting pyc file to source code: {self.file_path}') |
| 205 | sourcecode_text = subprocess.check_output( |
| 206 | [pycdc_path, self.file_path], stderr=subprocess.DEVNULL) |
| 207 | to_reconstruct = extract_broken_functions(sourcecode_text) |
| 208 | # reverse the list to start reconstructing from the bottom of the file, reducing the need to shift line numbers in the future |
| 209 | to_reconstruct = to_reconstruct[::-1] |
| 210 | |
| 211 | sourcecode_text_lines = sourcecode_text.split(b'\n') |
| 212 | sourcecode_text_lines = convert_indent_to_spaces(sourcecode_text_lines) |
| 213 | |
| 214 | model = Model() |
| 215 | for class_name, func_name, (start, end) in to_reconstruct: |
| 216 | logger.info( |
| 217 | f"Reconstructing function {func_name} in class {class_name}") |
| 218 | code_obj = get_function_from_code_objects( |
| 219 | file_as_code_object, class_name, func_name) |
| 220 | if not code_obj: |
| 221 | raise ValueError( |
| 222 | f"{self.file_path} : Function {func_name} in class {class_name} not found in code object") |
| 223 | |
| 224 | # in an effort to reduce token count, we will remove as much whitespace as possible from the bytecode |
| 225 | lowest_indent = min([len(x) - len(x.lstrip()) |
| 226 | for x in code_obj.disassembly]) |
| 227 | bytecode = '\n'.join([x[lowest_indent:] |
| 228 | for x in code_obj.disassembly]) |
| 229 | |