Обнаружение 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: Код, отправленный на проверку в компанию, предоставившую это тестовое задание, находится здесь.