import os
import pefile
import shutil
from tqdm import tqdm
def is_ELF_file(file_path):
    """
    判断样本是否是ELF文件
    :param file_path:样本文件路径
    :return:是ELF文件返回True，否则返回False
    """
    # 定义ELF文件的魔数
    ELF_MAGIC_NUMBER = b'\x7fELF'

    # 读取文件前几个字节
    with open(file_path, 'rb') as f:
        header = f.read(4)

    # 判断文件是否为ELF文件
    if header == ELF_MAGIC_NUMBER:
        return True
    else:
        return False


def is_pe_file(file_path):
    """
    判断样本是否是PE文件
    :param file_path:样本文件路径
    :return:是PE文件返回True，否则返回False
    """
    try:
        with open(file_path, 'rb') as f:
            header = f.read(4)
            pe = pefile.PE(file_path)
            return True
    except pefile.PEFormatError:
        return False
def importKnownPackers(filePath):
    knownPackers =[]
    try:
        with open(filePath) as packerPath:
            for line in packerPath:
                knownPackers.append(line.strip())
    except:
        print("Error reading given file.")
        return None

    return knownPackers


def detectPackers(pe, knownPackers):
    sections=[]
    for section in pe.sections:
        sections.append(section.Name.decode('utf-8').strip('\x00'))
    matches=[]

    for section in sections:
        for item_section in knownPackers:
            if section == item_section or item_section==item_section.swapcase():
                matches.append(section)

    return matches


def detect_pack_res(targetMalware):
    knownPackerFile = "knownPackerSections.txt"
    targetMalware = pefile.PE(targetMalware)
    knownPackers = importKnownPackers(knownPackerFile)
    detectedPackers = detectPackers(targetMalware, knownPackers)
    targetMalware.close()
    if detectedPackers:
        print("有加壳")
        return True
    else:
        print("无加壳")
        return False


class de_shell(object):
    def __init__(self,sample_path,save_dir):
        self.sample_path=sample_path
        self.save_dir=save_dir

    def fileFilter(self):
        for root, dirs, files in os.walk(self.sample_path):
            for file in tqdm(files):
                # 获取文件所属目录
                # 获取文件路径
                # 删除指定后缀名的文件
                targetMalwarepath = os.path.join(root, file)
                print(targetMalwarepath)
                try:
                    if detect_pack_res(targetMalwarepath):
                        shutil.copy(targetMalwarepath, self.save_dir)
                        os.remove(targetMalwarepath)
                        print('删除加壳文件' + file+'已保存在'+self.save_dir+'中，请查看')
                except Exception as e:
                    print(f"删除文件 {targetMalwarepath} 失败：{str(e)}")
        print("已留下清理后的样本文件")