Обнаружение DGA доменов или тестовое задание на позицию intern ML-engineer

В этой статье мы рассмотрим простую задачу, которая используется одной компанией в качестве тестового задания для стажеров на позицию ML-engineer. Она включает обнаружение DGA-доменов — задача, решаемая с помощью базовых инструментов машинного обучения. Мы покажем, как с ней справиться, применяя самые простые методы. Знание сложных алгоритмов важно, но куда важнее — понимать базовые концепции и уметь применять их на практике, чтобы успешно демонстрировать свои навыки.

DGA (Domain Generation Algorithm) — это алгоритм, который автоматически генерирует доменные имена, часто используемые злоумышленниками для обхода блокировок и связи с командными серверами.

В техничесокм задании присутствовали тестовые данные, для которых нужно было сформировать предсказания, и валидационные данные, на которых нужно было продемонстрировать метрики в формате:

True Positive (TP): False Positive (FP): False Negative (FN): True Negative (TN): Accuracy: Precision: Recall: F1 Score:

Иногда компании не предоставляют тренировочные данные и хотят оценить, насколько вы способны самостоятельно находить решения. Это включает:

1. Понимание проблемы: Четкое формулирование задачи.

2. Методология: Разработка плана действий и выбор методов.

3. Критическое мышление: Анализ данных и выдвижение гипотез.

4. Практические навыки: Применение базовых концепций машинного обучения.

Важно продемонстрировать инициативу и способность работать с ограниченной информацией. В нашем случае, домены существующих компаний можно найти на kaggle, а несуществующие домены нам необходимо будет сгенерировать самим.

Качественные и разнообразные данные позволяют алгоритмам выявлять закономерности, делать предсказания и принимать обоснованные решения. Поэтому без хороших данных невозможно достичь успешных результатов в машинном обучении. Важно создать качественные данные для обучения модели, чтобы обеспечить её эффективность и точность. Нам необходимо сосредоточиться на создании таких данных:

1. Напишем функции для генерации случайных строк и доменных имён. Функция generate_random_string генерирует строку заданной длины с буквами и, опционально, цифрами. Функция generate_domain_names создает список доменных имён с различными паттернами.

   def generate_random_string(length, use_digits=True):
     """
     Генерирует случайную строку заданной длины, включающую буквы и опционально цифры.
   
     :param length: Длина строки
     :param use_digits: Включать ли цифры в строку
     :return: Случайная строка
     """
     characters = string.ascii_lowercase
     if use_digits:
         characters += string.digits
     return ''.join(random.choice(characters) for _ in range(length)) 
   
   def generate_domain_names(count):
       """
       Генерирует список доменных имён с различными паттернами и TLD.
   
       :param count: Количество доменных имён для генерации
       :return: Список сгенерированных доменных имён
       """
       tlds = ['.com', '.ru', '.net', '.org', '.de', '.edu', '.gov', '.io', '.shop', '.co', '.nl', '.fr', '.space', '.online', '.top', '.info']
   
       def generate_domain_name():
           tld = random.choice(tlds)
           patterns = [
               lambda: generate_random_string(random.randint(5, 10), use_digits=False) + '-' + generate_random_string(random.randint(5, 10), use_digits=False),
               lambda: generate_random_string(random.randint(8, 12), use_digits=False),
               lambda: generate_random_string(random.randint(5, 7), use_digits=False) + '-' + generate_random_string(random.randint(2, 4), use_digits=True),
               lambda: generate_random_string(random.randint(4, 6), use_digits=False) + generate_random_string(random.randint(3, 5), use_digits=False),
               lambda: generate_random_string(random.randint(3, 5), use_digits=False) + '-' + generate_random_string(random.randint(3, 5), use_digits=False),
           ]
           domain_pattern = random.choice(patterns)
           return domain_pattern() + tld
   
       domain_list = [generate_domain_name() for _ in range(count)]
       return domain_list

2. Код загружает три CSV-файла, обрабатывает данные, удаляя столбец '1' и добавляя 'is_dga' со значением 0. Генерирует 1 миллион DGA-доменных имён, объединяет их с part_df и перемешивает итоговый DataFrame.

   try:
     logging.info('Загрузка данных')
     part_df = pd.read_csv('top-1m.csv')
     df_val = pd.read_csv('val_df.csv')
     df_test = pd.read_csv('test_df.csv')
     logging.info('Данные успешно загружены.')
   except Exception as e:
     logging.error(f'Ошибка при загрузке данных: {e}')
   
   logging.info('Обработка данных')
   part_df = part_df.drop('1', axis=1)
   part_df.rename(columns={'google.com': 'domain'}, inplace=True)
   part_df['is_dga'] = 0
   list_dga = df_val[df_val.is_dga == 1].domain.tolist()
   generated_domains = generate_domain_names(1000000)
   part_df_dga = pd.DataFrame({
       'domain': generated_domains,
       'is_dga': [1] * len(generated_domains)
   })
   df = pd.concat([part_df, part_df_dga], ignore_index=True)
   df = df.sample(frac=1).reset_index(drop=True)

3. Исключаем домены из валидационного и тестового наборов, затем балансируем классы, выбирая по 500,000 примеров для каждого из них. Итоговый сбалансированный набор перемешивается и сбрасывает индексы

   # Исключение доменов из валидационного и тестового наборов
   train_set = set(df.domain.tolist())
   val_set = set(df_val.domain.tolist())
   test_set = set(df_test.domain.tolist())
   intersection_val = train_set.intersection(val_set)
   intersection_test = train_set.intersection(test_set)
   if intersection_val or intersection_test:
     df = df[~df['domain'].isin(intersection_val | intersection_test)]
   
   
   # Балансировка классов до одинакового числа примеров
   logging.info('Балансировка классов')
   df_train_0 = df[df['is_dga'] == 0]
   df_train_1 = df[df['is_dga'] == 1]
   num_samples_per_class = 500000
   df_train_0_sampled = df_train_0.sample(n=num_samples_per_class, random_state=42)
   df_train_1_sampled = df_train_1.sample(n=num_samples_per_class, random_state=42)
   df_balanced = pd.concat([df_train_0_sampled, df_train_1_sampled])
   df_train = df_balanced.sample(frac=1, random_state=42).reset_index(drop=True)
   

4. Создаем и обучаем модель, используя конвейер, который включает векторизацию с помощью TfidfVectorizer и логистическую регрессию. После обучения модель сохраняется в файл model_pipeline.pkl

   logging.info('Создание и обучение модели')
   
   model_pipeline = Pipeline([
       ("vectorizer", TfidfVectorizer(tokenizer=n_grams, token_pattern=None)),
       ("model", LogisticRegression(solver='saga', n_jobs=-1, random_state=12345))
   ])
   
   model_pipeline.fit(df_train['domain'], df_train['is_dga'])
   logging.info('Сохранение модели')
   joblib_file = "model_pipeline.pkl"
   joblib.dump(model_pipeline, joblib_file)
   logging.info(f'Модель сохранена в {joblib_file}')

Вся наша задача сводится к тому, что нам необходимо домены разбить на N-граммы и векторизовать их с помощью TF-IDF. N-грамма — это последовательность из N элементов (слов или символов) в тексте, но в нашей задаче мы применяем их к одному слову, чтобы выделять и анализировать слоги доменов. TF-IDF (Term Frequency-Inverse Document Frequency) — это метод, который помогает оценить важность слова в документе по сравнению с другими документами в коллекции.

Таким образом, комбинируя N-граммы и TF-IDF, мы можем эффективно анализировать домены и выявлять их ключевые характеристики. Рассмотрим на примере существующих доменов: texosmotr-auto.ru и pokerdomru.ru, разобьем их на 4-граммы, не беря во внимание родовой домен (.ru)

• Для texosmotr-auto.ru: "texo", "exos", "xosm", "osmo", "smot", "motr", "otr-", "r-au", "-aut", "auto"

• Для pokerdomru.ru: "poke", "oker", "kerd", "erdo", "domr", "omru"

Мы рассмотрели 4-граммы, но разве для всех доменов необходимо использовать фиксированные N-граммы? Конечно, нет. Для каждого домена создаются 3-мерные, 4-мерные и 5-мерные граммы, чтобы выявить различные языковые паттерны и особенности структуры. Такой подход позволяет лучше захватывать контекст и увеличивает возможность обнаружения уникальных характеристик, которые могут быть полезны для классификации.

• код для созднания 3-мерных, 4-мерных и 5-мерных грамм для домена

  def n_grams(domain):
    """
    Генерирует n-граммы для доменного имени.
  
    :param domain: Доменное имя
    :return: Список n-грамм
    """
    grams_list = []
    # Размеры n-грамм
    n = [3, 4, 5]
    domain = domain.split('.')[0]
    for count_n in n:
        for i in range(len(domain)):
            if len(domain[i: count_n + i]) == count_n:
                grams_list.append(domain[i: count_n + i])
    return grams_list

Все полученные N-граммы необходимо векторизовать, и в этом нам поможет вышеупомянутый метод TF-IDF. Этот подход позволяет оценить важность каждой N-граммы в контексте доменов, преобразуя текстовые данные в числовую форму. Векторизация с помощью TF-IDF учитывает частоту встречаемости N-грамм в каждом домене и их редкость в общем наборе.

Финальным этапом необходимо обучить нашу модель. Вы можете использовать различные алгоритмы, которые улучшают вашу метрику, однако я выбрал классическую логистическую регрессию (LR), потому что она проста в реализации, хорошо интерпретируется и часто дает неплохие результаты, например я получил следующие метрики на валидационном наборе данных:

True Positive (TP): 4605 False Positive (FP): 479 False Negative (FN): 413 True Negative (TN): 4503 Accuracy: 0.9108 Precision: 0.9058 Recall: 0.9177 F1 Score: 0.9117

Таким образом, понимание базовых концепций, таких как N-граммы и TF-IDF, откроет перед вами возможности для решения прикладных задач и позволит уверенно заявить о себе на стажировках. Эти навыки станут крепкой основой для вашего профессионального роста в области машинного обучения и анализа данных.

PS: Код, отправленный на проверку в компанию, предоставившую это тестовое задание, находится здесь.