from oletools.olevba import VBA_Parser, filter_vba from tqdm import tqdm from typing import List import logging import traceback import pickle import argparse import os import sys __version__ = '1.0.1' PROGRESS_FILE = 'progress.pkl' OUTPUT_FOLDERNAME = 'results' logging.basicConfig( level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler('error.log', encoding='utf-8'), ] ) def parse_args(args: List[str] = sys.argv[1:]): parser = argparse.ArgumentParser(description='Скрипт предназначен для извлечения vba макросов из документов формата xls, xlsm.') parser.add_argument('--file', '-f', type=str, required=True, help='Путь к текстовому файлу со списком документов.') return parser.parse_args(args) def processing_files(list_files): try: if not list_files: print('Список файлов пуст') sys.exit(1) current_progress = check_progress() if current_progress == len(list_files): print('Все файлы уже обработаны, удалите файл прогресса и запустите скрипт заново.') sys.exit(1) try: for index in tqdm(range(current_progress, len(list_files)), desc='Обработка файлов'): if not os.path.exists(list_files[index]): raise FileNotFoundError(f'File not found - {list_files[index]}') if not check_file_access(list_files[index]): raise PermissionError(f'Permission denied - {list_files[index]}') dirname = check_folder_existence(generate_dirname(list_files[index])) extract_macro_from_doc(dirname, list_files[index]) display_processed_files_count(index, len(list_files)) delete_progress_file(PROGRESS_FILE) except (FileNotFoundError, PermissionError) as e: save_progress(index) tb = get_traceback(e) display_processed_files_count(index - 1, len(list_files)) print_error(e, index, tb) except ValueError as e: save_progress(index) tb = get_traceback(e) display_processed_files_count(index - 1, len(list_files)) print_error(e, index, tb) except Exception as e: save_progress(index) tb = get_traceback(e) display_processed_files_count(index - 1, len(list_files)) print_error(e, index, tb) def processing_single_file(path_to_file): try: if not os.path.exists(path_to_file): raise FileNotFoundError(f'Файл "{path_to_file}" не существует.') else: if check_extension(path_to_file): dirname = check_folder_existence(generate_dirname(path_to_file)) extract_macro_from_doc(dirname, path_to_file) else: raise ValueError(f'Файл "{path_to_file}" имеет некорректное расширение. Ожидается xls или xlsm.') except FileNotFoundError as e: print(f'Ошибка: {e}') except ValueError as e: print(f'Ошибка: {e}') def display_processed_files_count(current_step, total_steps): print(f'Обработано файлов: {current_step + 1}/{total_steps}') def get_traceback(error): tb = traceback.extract_tb(error.__traceback__)[-1] return tb def check_file_access(filepath): try: with open(filepath, 'rb'): pass return True except Exception: return False def check_folder_existence(dirname): if os.path.exists(OUTPUT_FOLDERNAME): dirname = os.path.join(OUTPUT_FOLDERNAME, dirname) if os.path.exists(dirname): dirname = create_folder_with_suffix(dirname) else: os.mkdir(dirname) else: dirname = os.path.join(OUTPUT_FOLDERNAME, dirname) os.makedirs(dirname, exist_ok=True) return dirname def check_progress(): if not load_progress(): current_progress = 0 else: current_progress = load_progress() return current_progress def extract_macro_from_doc(dirname, doc): for _, _, filename, content in get_vba_modules(doc): filtered_content = filter_vba(content) path_to_file = os.path.join(dirname, os.path.join('.', filename)) generate_files(path_to_file, filtered_content) def get_vba_modules(workbook_path): vba_parser = VBA_Parser(workbook_path) vba_modules = vba_parser.extract_all_macros() if vba_parser.detect_vba_macros() else [] return vba_modules def generate_files(path_to_file, filtered_content): with open(path_to_file, 'w', encoding='utf-8') as f: f.write(filtered_content) def generate_dirname(path_to_workbook): if path_to_workbook: dirname = os.path.splitext(os.path.basename(path_to_workbook))[0] if check_extension(path_to_workbook): return dirname else: return False else: return False def check_extension(path_to_workbook): file_extension = os.path.splitext(os.path.basename(path_to_workbook))[1] if file_extension == '.xls' or file_extension == '.xlsm': return True else: return def get_list_of_files(folder_path): file_names = [] for file_name in os.listdir(folder_path): if check_startswith(os.path.join(folder_path, file_name)): file_names.append(os.path.join(folder_path, file_name)) return file_names def check_startswith(path_to_workbook): file_name = os.path.basename(path_to_workbook) if ( os.path.isfile(path_to_workbook) and check_extension(file_name) and not file_name.startswith(('.', '~', '$')) ): return True else: return def get_list_of_folders(folder_path): dir_names = [] for dir_name in os.listdir(folder_path): if ( os.path.isdir(os.path.join(folder_path, dir_name)) and not dir_name.startswith('.') ): dir_names.append(os.path.join(folder_path, dir_name)) return dir_names def create_folder_with_suffix(base_folder_name): folder_name = base_folder_name counter = 1 while os.path.exists(folder_name): folder_name = f'{base_folder_name}_{counter}' counter += 1 os.mkdir(folder_name) return folder_name def save_progress(progress_data): with open(PROGRESS_FILE, 'wb') as f: pickle.dump(progress_data, f) def load_progress(): if os.path.exists(PROGRESS_FILE): with open(PROGRESS_FILE, 'rb') as f: progress_data = pickle.load(f) return progress_data else: return def delete_progress_file(progress_file): if os.path.exists(progress_file): os.remove(progress_file) else: return def get_all_files_recursive(path_to_dir): files_list = [] for root, dirs, files in os.walk(path_to_dir): for file in files: if check_extension(os.path.join(root, file)) and check_startswith(os.path.join(root, file)): files_list.append(os.path.join(root, file)) return files_list def get_docs_path_from_file(file_path): with open(file_path, 'r', encoding='utf-8') as file: return [line.strip() for line in file if line.strip()] def get_list_with_all_files(list_files_from_txt): processed_files = set(list_files_from_txt) for file in tqdm(list_files_from_txt, desc='Подготовка файлов'): folder_path = os.path.dirname(file) all_file_recursive = get_all_files_recursive(folder_path) processed_files.update(all_file_recursive) processed_files = sorted(processed_files) return processed_files def print_error(error, index, tb_info=None): first_message = f'Ошибка: {error}' last_message = f'Сохранение прогресса и завершение работы...' if tb_info: detailed_message = ( f'Функция: {tb_info.name}\n' f'Строка: {tb_info.lineno}\n' ) full_message = f'{first_message}\n{detailed_message}{last_message}' logging.error(f'{tb_info.filename} | {tb_info.lineno} | document index: {index} | {tb_info.name} | {error}') else: full_message = f'{first_message}\n{last_message}' print(full_message) def main(): arguments = parse_args() if arguments.file: list_files = get_docs_path_from_file(arguments.file) processing_files(list_files) else: print( ''' usage: extract_vba.py [-h] --file FILE extract_vba.py: error: the following arguments are required: --file ''' ) if __name__ == '__main__': main()