Пишем свой антивирус с нуля часть 2

Март 25, 2024, 15:03 - Время чтения: 2 минуты

Как и обещал в предыдущей статье - вот простейший сканер:

import os
import hashlib
from sqlalchemy import create_engine, Column, String
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

db_path = "file_hashes.db"

engine = create_engine(f"sqlite:///{db_path}")
Session = sessionmaker(bind=engine)
session = Session()

Base = declarative_base()

class FileHash(Base):
    __tablename__ = "file_hashes"
    id = Column(String, primary_key=True)
    filename = Column(String)
    sha256 = Column(String)

Base.metadata.create_all(engine)

def calculate_sha256(file_path):
    hasher = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hasher.update(chunk)
    return hasher.hexdigest()

def add_file_to_db(file_path):
    filename = os.path.basename(file_path)
    sha256 = calculate_sha256(file_path)

    existing_hash = session.query(FileHash).filter_by(sha256=sha256).first()

    if existing_hash:
        print(f"Файл {file_path} уже существует в базе данных.")
        return
    file_hash = FileHash(id=sha256, filename=filename, sha256=sha256)
    session.add(file_hash)
    session.commit()
    print(f"Файл {file_path} успешно добавлен в базу данных.")

def scan_directory(directory):
    n = 0
    v = 0
    for root, dirs, files in os.walk(directory):
        for filename in files:
            file_path = os.path.join(root, filename)
            sha256 = calculate_sha256(file_path)

            existing_hash = session.query(FileHash).filter_by(sha256=sha256).first()
            n += 1
            if existing_hash:
                v += 1
                print(f"{file_path}: Вероятно вирус!")
    return n, v

if __name__ == "__main__":
    directory_to_scan = "/test/"
    n, v = scan_directory(directory_to_scan)

    print(f"Просканировано: {n} файлов")
    print(f"Найдено: {v} угроз")

Для теста написал неприятную программку которая сильно кушает процессорное время:

#include <stdlib.h>

int main() {
    for (;;) { fork(); free(malloc(4096)); }
    return -3;
}

Лучше не запускать её, но я скомпилировал и проверил :) Теперь нужно добавить в базу это недоразумение:

if __name__ == "__main__":
    directory_to_scan = "/test/"
    add_file_to_db('/test/a.out')
    n, v = scan_directory(directory_to_scan)

    print(f"Просканировано: {n} файлов")
    print(f"Найдено: {v} угроз")

Вот примерный вывод программы:

Файл /test/a.out успешно добавлен в базу данных.
/test/a.out: Вероятно вирус!
Просканировано: 4 файлов
Найдено: 1 угроз

В следующей статье будет более наглядный и правильный пример антивирусной программы