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

Комментарии (0)