Pull to refresh

Пароли в открытом доступе: ищем с помощью машинного обучения

Level of difficultyMedium
Reading time18 min
Views9.7K

Я больше 10 лет работаю в сфере IT и информационной безопасности. И знаю, что сложнее всего предотвратить риски, связанные с человеческим фактором. 

Мы разрабатываем самые надежные способы защиты. Но всего один оставленный в открытом доступе пароль сведет все усилия к нулю. А чего только не отыщешь в тикетах и комментариях в Jira, правда?

Привет, меня зовут Саша Рахманный, я разработчик в команде информационной безопасности в Lamoda Tech. В этой статье поделюсь опытом, как мы ищем в корпоративных ресурсах чувствительные данные — пароли, токены и строки подключения, — используя самописный ML-плагин. Рассказывать о его реализации буду по шагам и с подробностями, чтобы вы могли создать такой инструмент у себя, даже если ML для вас — незнакомая технология.  

Как пароли оказываются в открытом доступе 

Как-то раз на одной из моих прошлых работ мы проводили рядовой тест на проникновение. Результаты были удручающими: пентестер не просто проник в систему, но и получил доступ ко всем учетным записям. Слабым местом оказался пароль, открыто опубликованный в тикете Jira.

Пентестер действовал таким образом:

  1. Обманом получил пользовательские данные для входа в Service Desk компании. Пользователь, чьи данные использовали, не был сотрудником ИТ-подразделения.

  2. С помощью поиска и regexp нашел тикет, в котором администратора домена просили создать сервисную учетную запись в Active Directory и сделать её локальным администратором на одном из серверов. В ответ администратор указал в комментарии логин и пароль от учетной записи.

  3. Пентестер подключился по RDP к серверу с этими реквизитами и обнаружил, что администратор не завершил сеанс. 

  4. С помощью обфусцированной версии Mimikatz получил пароль администратора домена. Установленный антивирус Касперского не обнаружил ничего подозрительного.

  5. Пентестер сдампил хеши паролей всех учетных записей в домене. На этом тест был завершен.

Конечно, мы разобрали ситуацию и приняли все нужные меры. Закрыли возможность просматривать проекты Service Desk, ограничили срок сессии, организовали подключение привилегированных учеток с использованием usb-токенов и так далее. 

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

Прошло несколько лет, я уже работал в Lamoda в отделе информационной безопасности. Мне попался на глаза тикет, в котором была проблема с подключением к к одной из систем компании с токеном в открытом виде.

После пережитого флешбека к тому пентесту стало интересно, а сможем ли мы в реальном времени узнавать о таких инцидентах?  

Ищем ML-решение

В первую очередь мы начали искать готовое решение. Но оказалось, что готовых opensource-плагинов для Jira на этот случай нет. Поэтому мы решили написать свой с использованием Machine Learning и сделать его более универсальным, для подключения к различным источникам: Jira, Confluence, общие диски. Создать некий аналог DLP, но для технических данных. 

Одно из существенных ограничений — модель должна работать локально, никакие данные во внешний мир мы передавать не должны.

В качестве фреймворка была выбрана ML.NET — opensource-библиотека Microsoft для машинного обучения для .NET приложений. 

Почему именно .NET, а не Python? Тут все просто: технический стек команды ИБ в Lamoda Tech — это продукты на .NET, и мы выбрали то, что знали и с чем работали. Так же фреймворк ML.NET позволяет обучать модели без глубокого погружения в машинное обучение благодаря конструкторам моделей.

Предлагаю в статье пройти весь путь от сбора данных для обучения до публикации готового приложения. 

Собираем данные для обучения модели

В самом начале нам нужны тренировочные данные для обучения. Для удобства разработки, обучения и оценки мы на первом этапе разбили все типы данных на три категории: пароли, токены и строки подключения. 

В качестве источника выбрали корпоративная Jira с многолетней историей тикетов. Алгоритм действий был такой:

Создали набор регулярных выражений

За основу был взят проект https://github.com/mazen160/secrets-patterns-db. Мы исключили те системы из списка, которые не используем и не планируем использовать, — данных для обучения мы по ним не получим. Добавили специфичные для нас выражения: например, вариации слова «пароль» в латинице, токены местных провайдеров (Yandex Cloud), регулярные выражения для Bearer и basic-токенов.

Выгрузили тикеты из Jira по API, используя поиск по ключевым словам

К сожалению, встроенный язык запросов Jira JQL не поддерживает поиск по регулярным выражениям. Мы составили список ключевых слов, в которых могут встречаться искомые данные, и выгрузили все тикеты, соответствующие им.

Создали CSV-файлы под каждую категорию

Разбили содержимое тикетов — описание, комментарии — на отдельные строки и проверили каждую на соответствие регулярному выражению. Записали результат в три файла: строки подключения, токены, пароли. Данные, соответствующие регулярному выражению, пометили как True, остальные данные — False. Пример CSV файла:

Text

ContainsToken

PASSWORD=abc123xyz

True

ALLOWED_HOSTS=example.com

False

где столбец Text — текст для обучения, ContainsToken — True, если содержит, или False, если нет.

Обучаем модели

Для обучения и работе с моделями создадим отдельную библиотеку Net Core, модели из которой мы потом сможем использовать в Web API сервере. 

Процесс создания модели достаточно прост.

Необходимо в Visual Studio перейти в проект библиотеки, Add — Machine Learning Model. Задать имя модели.

После чего откроется мастер создания модели. Подробнее о сценариях ML.NET можно почитать тут

В нашем сценарии нас интересует Data classification, мы будем обучать модель отвечать на вопрос, присутствует ли во входящей строке пароль (секрет/токен)? На первом этапе модель у нас выступает в роли бинарного классификатора. 

В качестве данных для обучения указываем подготовленный на предыдущем этапе CSV-файл. В настройках указываем, что колонка Text — это строковые данные для обучения, а ContainsToken выступает в качестве категории.

В меню Validation data в конструкторе можно выбрать три стратегии:

  1. Cross validation — это метод обучения и оценки моделей машинного обучения, который разбивает данные на несколько частей и обучает несколько алгоритмов на этих частях. Этот метод рекомендуется применять при небольшом количестве данных для обучения.

  2. Split — предоставленные данные разделяются на обучение и проверку (например, 80/20). Этот алгоритм стоит применять при большом количестве исходных данных.

  3. Validation data — позволяет указать отдельный файл для проверки модели.

Мы сравнивали стратегии Cross validation и Split. Валидация Cross validation хоть и дольше, но показывает лучшие результаты на тех же тестовых данных. 

Пример:

Строка для теста

curl -X GET "https://api.somesite.com/api/someversion/someendpoint" \ -H "Content-Type:application/json" \ -H "Authorization: Bearer 123456789ABCDEF"

Результат:

Cross validation, 5 folds  Score: 4.301

Split, 80/20, Score: 3.725

Далее нужно настроить обучение модели.

Настройки обучения модели

Все настройки, кроме Optimizing metric, появились с версии ML.NET 3, поэтому лучше использовать последнюю версию фреймворка.

Time to train — можно ограничить время обучения модели в секундах. Значение зависит и от количества тестовых данных и от производительности процессора. Microsoft рекомендует придерживаться следующих ограничений:

На моем датасете в 5 Мб не было существенной разницы между 10 и 60 минутами.

Advanced — Optimizing metric — это специфические показатели, которые используются для оценки качества модели машинного обучения, выполняющей определенную задачу. 

Типы метрик:

  1. Accuracy — это доля правильных прогнозов с набором тестовых данных. Это отношение количества правильных предсказаний к общему количеству входных выборок.

  2. AUC-ROC — площадь под кривой ошибок (ROC), которая показывает зависимость между долей истинно положительных и ложно положительных предсказаний модели при разных порогах классификации. Чем выше AUC-ROC, тем лучше модель разделяет классы.

  3. AUCPR — кривая точности-полноты (PR-curve) показывает зависимость между точностью (precision) и полнотой (recall) модели при разных порогах классификации. 

Точность — это доля правильных предсказаний модели среди объектов, которые модель отнесла к положительному классу. 

Полнота — это доля правильных предсказаний модели среди объектов, которые действительно принадлежат положительному классу. Площадь под кривой точности-полноты (PR-AUC) — это среднее значение точности, полученное для каждого значения полноты.

  1. F1-score — F1-мера, также известная как сбалансированная F-мера или F-оценка, — это среднее гармоническое между точностью и полнотой. F1-мера полезна, когда вы хотите найти баланс между точностью и полнотой. Математическая формула для F1-меры следующая: F1 = 2 * (precision * recall) / (precision + recall).

На наших данных для обучения особой разницы между метриками обучения мы не заметили, используем Accuracy. Подробнее о метриках можно прочитать здесь: towardsdatascience.com.

Advanced — Trainers — алгоритм обучения. В первый раз можно выбрать все алгоритмы и проверить, какой из них показывает лучший результат. В ML.NET 2 лучший результат на наших данных показывал TextClassificationMulti (0.9756), в ML.NET 3 на тех же данных LightGbmBinary (0.9961). Подробнее о алгоритмах обучения читайте здесь: learn.microsoft.com.

Advanced — Tuners — алгоритмы перебора параметров:

  1. Stop strategy (Остановка обучения по заданным параметрам) — время, количество моделей и лимит оперативной памяти.

  2. Simple strategy — использование только части датасета для обучения. 

После обучения конструктор покажет результат лучшей модели:

На последнем этапе конструктор даст пример кода для использования модели либо предложит создать проект с консольным приложением или Web API. На этом этапе мы закрываем конструктор и создаем аналогичные модели для других типов — паролей и строк подключения.

Проверка работы моделей, API-сервер

Для проверки работы моделей создадим WebAPI-приложение с единственным эндпоинтом — Predict. Для демонстрации работы на данном этапе упрощена работа с PredictionEngine.

Для работы сервиса создали интерфейс:

public interface IPredictor
{
    Task<PredictResult> PredictAsync(string str);
}

И классы, реализующие его.

Класс для проверки по регулярным выражениям:

public class RegexPredictor : IPredictor
{
    private readonly RegexConfig _regexConfig;
    public RegexPredictor()
    {
        var deserializer = new YamlDotNet.Serialization.Deserializer();
        _regexConfig = deserializer.Deserialize<RegexConfig>(File.ReadAllText("rules.yml"));
    }

    /// <summary>
    /// predict using regex
    /// </summary>
    /// <param name="str"></param>
    /// <returns>PredictResult</returns>
    public Task<PredictResult> PredictAsync(string str)
    {
       
        foreach (var item in _regexConfig.patterns)
        {
            Match match;
            try
            {
                match = Regex.Match(str, item.pattern.regex, RegexOptions.IgnoreCase);
            }
            catch
            {
                 match = Regex.Match(str, Regex.Escape(item.pattern.regex), RegexOptions.IgnoreCase);
            }
            if (match.Success)
            {
                return Task.FromResult(new PredictResult()
                {
                    Name = $"Regex({item.pattern.name})",
                    Score = 1

                });
            }
        }
        return Task.FromResult(new PredictResult()
        {
            Score = -1,
            Name = "Regex"

        });
    }
}

Абстрактный класс для оценки модели:

public abstract class MLPredictor
{

    /// <summary>
    /// Create a prediction engine from a ML.NET model
    /// </summary>
    /// <param name="mlNetModelPath"></param>
    /// <returns>PredictionEngine</returns>
    internal  PredictionEngine<ModelInput, ModelOutput> CreatePredictEngine(string mlNetModelPath)
    {
        var mlContext = new MLContext();
        ITransformer mlModel = mlContext.Model.Load(mlNetModelPath, out var _);
        return mlContext.Model.CreatePredictionEngine<ModelInput, ModelOutput>(mlModel);
    }

    /// <summary>
    /// Predict the input string
    /// </summary>
    /// <param name="input"></param>
    /// <param name="predictEngine"></param>
    /// <returns>ModelOutput</returns>
    public ModelOutput Predict(ModelInput input, Lazy<PredictionEngine<ModelInput, ModelOutput>> predictEngine)
    {
        var predEngine = predictEngine.Value;
        return predEngine.Predict(input);
    }
}

И классы реализации для каждой модели, на примере модели с паролями:

public class PasswordPredictor : MLPredictor, IPredictor
{
    private readonly Lazy<PredictionEngine<ModelInput, ModelOutput>> _predictEnginePassword;

    public PasswordPredictor()
    {
        _predictEnginePassword = new Lazy<PredictionEngine<ModelInput, ModelOutput>>(() => 
            CreatePredictEngine("PasswordModel.mlnet"), true);
    }

    /// <summary>
    /// Predict the input string
    /// </summary>
    /// <param name="str"></param>
    /// <returns></returns>
    public async Task<PredictResult> PredictAsync(string str)
    {
        var predictResult = Predict(new ModelInput
        {
            Text = str
        }, _predictEnginePassword);
        return new PredictResult
        {
            Name = "Password",
            Score = predictResult.Score
        };

    }
}

Общий класс для вывода лучшего результата:

public class CombinePredictor : IPredictor
{
    private readonly List<IPredictor> _predictors;

    public CombinePredictor(List<IPredictor> predictors)
    {
        _predictors = predictors;
    }

    /// <summary>
    /// Run all predictors and return the result with highest confidence
    /// </summary>
    /// <param name="str"></param>
    /// <returns>PredictResult</returns>
    public async Task<PredictResult> PredictAsync(string str)
    {
        var tasks = _predictors.Select(p => p.PredictAsync(str));
        var results = await Task.WhenAll(tasks);
        var max = results.Max(r => r.Score);
        var result = results.First(r => r.Score == max);
        return Task.FromResult(result);
    }
}

И класс для возврата результата:

/// <summary>
/// Predict Result model
/// </summary>
public class PredictResult
{
    /// <summary>
    /// Name of the predictor
    /// </summary>
    public string? Name { get; set; }
    /// <summary>
    /// Score of the prediction
    /// </summary>
    public float Score { get; set; }
    /// <summary>
    /// Is credential?
    /// </summary>
    public bool IsTrue => Score > 0;
}

Стандартные модели из ML.NET:

public class ModelInput
{
    [LoadColumn(0)]
    [ColumnName(@"Text")]
    public string Text { get; set; }

    [LoadColumn(1)]
    [ColumnName(@"ContainsToken")]
    public bool ContainsToken { get; set; }

}
public class ModelOutput
{
    [ColumnName(@"Text")]
    public float[] Text { get; set; }

    [ColumnName(@"ContainsToken")]
    public bool ContainsToken { get; set; }

    [ColumnName(@"Features")]
    public float[] Features { get; set; }

    [ColumnName(@"PredictedLabel")]
    public bool PredictedLabel { get; set; }

    [ColumnName(@"Score")]
    public float Score { get; set; }

    [ColumnName(@"Probability")]
    public float Probability { get; set; }

}

После чего необходимо зарегистрировать наши сервисы моделей и создать API эндпоинт, для чего изменим Program.cs:

using WebApplication1.Services;

var builder = WebApplication.CreateBuilder(args);

builder.Services.AddEndpointsApiExplorer();
builder.Services.AddSwaggerGen();
builder.Services.AddSingleton<RegexPredictor, RegexPredictor>();
builder.Services.AddSingleton<PasswordPredictor, PasswordPredictor>();
builder.Services.AddSingleton<TokenPredictor, TokenPredictor>();

var app = builder.Build();
app.UseSwagger();
app.UseSwaggerUI();
app.MapGet("/predictCred", async (RegexPredictor regexPredictor, PasswordPredictor passwordPredictor, TokenPredictor tokenPredictor, string str) =>
{
    var predictors = new List<IPredictor>
    {
        regexPredictor,
        passwordPredictor,
        tokenPredictor
    };
    var combinePredictor = new CombinePredictor(predictors);
    var result = await combinePredictor.PredictAsync(str);
    return result;
}).WithName("Predict string");
app.Run();

После того, как все готово, можно запустить и проверить работу:

Интерпретация результата

Проверка по регулярным выражениям возвращает результат:

  •  -1, если не подошло ни одно выражение, 

  • 1, если подошло хотя бы одно.

ML.NET возвращает значения в гораздо большем диапазоне, в CombinePredictor выбирается лучший результат и возвращается в API.

На этом этапе можно подключить к API парсер, который будет просматривать ресурсы и проверять, есть ли чувствительные данные. На первом этапе мы использовали Jira: сервис по API запрашивал новые и измененные тикеты за 5 минут. Если реквизиты были найдены, оповещал нас в корпоративном мессенджере. В БД мы записываем MD5-хэш строки и тикет, в которых найдены реквизиты (для исключения повторного уведомления).

Дообучаем модели

После деплоя нашего сервиса мы можем столкнуться с некоторым количеством ложноположительных или ложноотрицательных оценок моделью. Для минимизации необходимо дообучить модель. 

Для этого нам необходимо сохранять запросы, результат и время, потраченное на оценку (для статистики). Для этого расширим модель, добавив свойство Duration:

/// <summary>
/// Predict Result model
/// </summary>
public class PredictResult
{
    /// <summary>
    /// Name of the predictor
    /// </summary>
    public string? Name { get; set; }
    /// <summary>
    /// Score of the prediction
    /// </summary>
    public float Score { get; set; }
    /// <summary>
    /// Is credential?
    /// </summary>
    public bool IsTrue => Score > 0;
    /// <summary>
    /// Duration of the prediction
    /// </summary>
    public int Duration { get; set; }
}

В каждом классе добавим таймер:

public class PasswordPredictor : MLPredictor, IPredictor
{
    private readonly Lazy<PredictionEngine<ModelInput, ModelOutput>> _predictEnginePassword;
    private readonly Stopwatch _timer = new Stopwatch();
    public PasswordPredictor()
    {
        _predictEnginePassword = new Lazy<PredictionEngine<ModelInput, ModelOutput>>(() => 
            CreatePredictEngine("PasswordModel.mlnet"), true);

    }
    /// <summary>
    /// Predict the input string
    /// </summary>
    /// <param name="str"></param>
    /// <returns></returns>
    public async Task<PredictResult> PredictAsync(string str)
    {
        _timer.Reset();       
        _timer.Start();
        var predictResult = Predict(new ModelInput
        {
            Text = str
        }, _predictEnginePassword);
        _timer.Stop();
        return new PredictResult
        {
            Name = "Password",
            Score = predictResult.Score,
            Duration = (int)_timer.ElapsedMilliseconds
        };

    }
}

Создаем сервис для дампа в CSV-файл:

/// <summary>
/// Dumper
/// </summary>
public class Dumper
{
    private const string FileName = "result.csv";
    /// <summary>
    /// Dump prediction result
    /// </summary>
    /// <param name="baseString"></param>
    /// <param name="results"></param>
    /// <returns></returns>
    internal static void Dump(string baseString,IEnumerable<PredictResult> results)
    {
        var sb = new StringBuilder();
        sb.Append($"{baseString};");
        foreach (var result in results)
        {
            sb.Append($"{result.Name};{result.Score};{result.IsTrue};{result.Duration};");
        }
        sb.AppendLine();
        File.AppendAllText(FileName, sb.ToString(), Encoding.UTF8);
    }
}

И добавляем дампер в основной класс:

public class CombinePredictor : IPredictor
{
    private readonly List<IPredictor> _predictors;
    private readonly Stopwatch _timer = new Stopwatch();
    public CombinePredictor(List<IPredictor> predictors)
    {
        _predictors = predictors;
    }
    /// <summary>
    /// Run all predictors and return the result with highest confidence
    /// </summary>
    /// <param name="str"></param>
    /// <returns>PredictResult</returns>
    public async Task<PredictResult> PredictAsync(string str)
    {
        _timer.Reset();
        _timer.Start();
        var tasks = _predictors.Select(p => p.PredictAsync(str));
        var results = await Task.WhenAll(tasks);
        var max = results.Max(r => r.Score);
        var result = results.First(r => r.Score == max);
        _timer.Stop();
        result.Duration = (int)_timer.ElapsedMilliseconds;
        Dumper.Dump(str, results);
        return Task.FromResult(result);
    }
}

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

Дальше необходимо прогнать через сервис большое количество данных для статистического анализа расхождений. Для этого был выбран массив тикетов за последние 6 месяцев (200К тикетов и десятки миллионов строк), но нас интересуют только расхождения между регулярными выражениями и машинным обучением. 

Для этого созданный CSV-файл необходимо проанализировать на наличие ложных срабатываний. Создадим класс:

public class Result
{
    [Index(0)]
    public string Str { get; set; }
    [Index(1)]
    public string NameRegex { get; set; }
    [Index(2)]
    public bool IsTrueRegex { get; set; }
    [Index(3)]
    public float ScoreRegex { get; set; }
    [Index(4)]
    public int DurationRegex { get; set; }


    [Index(5)]
    public string NamePassword { get; set; }
    [Index(6)]
    public bool IsTruePassword { get; set; }
    [Index(7)]
    public float ScorePassword { get; set; }
    [Index(8)]
    public int DurationPassword { get; set; }

    [Index(9)]
    public string NameToken { get; set; }
    [Index(10)]
    public bool IsTrueToken { get; set; }
    [Index(11)]
    public float ScoreToken { get; set; }
    [Index(12)]
    public int Duration { get; set; }
    public override string ToString()
    {
        return
            $"{Str} {NameRegex} {IsTrueRegex}, Password: {IsTruePassword}({ScorePassword}), Token: {IsTrueToken}({ScoreToken})";
    }
}

И новое консольное приложение:

var config = new CsvConfiguration(CultureInfo.CurrentCulture)
{
    HasHeaderRecord = false,
    Delimiter = ";",
    BadDataFound = null
};
var reader = new StreamReader(@"result.csv");
var csv = new CsvReader(reader, config);
var records = csv.GetRecords<Result>().ToList();
foreach (var record in records)
{
    if (record.IsTrueRegex && (!record.IsTruePassword && !record.IsTrueToken) || (!record.IsTrueRegex && (record.IsTruePassword || record.IsTrueToken)))
    {
        Console.WriteLine(record.ToString());
    }
}

На основании разницы необходимо провести анализ, дополнить тренировочные данные, провести новое обучение модели. Таких итераций необходимо несколько, с использованием разных данных для сравнения. Например, мы провели 8 итераций за 8 лет тикетов. 

После нескольких дообучений на выходе мы должны получить модели, которые предсказывают не хуже, чем регулярные выражения. В идеальной ситуации модели показывают даже лучшие результаты, так как в тексте могут быть, например, лишние пробелы или опечатки в слове «пароль», которые не найдут регулярные выражения, а ML, хоть и с более низкой оценкой, обнаружит.

Консолидируем модели в одну

После того, как мы подготовили оптимальные данные для обучения, пришло время для единой модели. Для этого необходимо все данные для обучения свести в единый файл. В качестве метки в нем будет выступать не True/False, а тип данных: Password, Secret, Token, ConnectionString, Clean. 

Примечание: необходимо обработать ситуации, когда в данных для обучения модели паролей строки токенов были помечены как False. При сведении в один файл лучше пройти по нему регулярными выражениями для проверки, иначе модель будет менее точная.

После сведения в файл создаем новую модель. Аналогично второму этапу выбираем Data Classification, но в Column settings указываем, что поле ContainsToken имеет тип String.

Обучение такой модели требует гораздо большего времени, в отличие от бинарного классификатора. Рекомендую установить не менее 30 минут для оценки работы как можно большего количества моделей. Также результат анализа может незначительно снизится.

После успешного обучения модели добавляем в проект API сервера по аналогии с другими, меняем модели ModelInput: 

public class ModelInput
{
    [LoadColumn(0)]
    [ColumnName(@"Text")]
    public string Text { get; set; }

    [LoadColumn(1)]
    [ColumnName(@"ContainsToken")]
    public string ContainsToken { get; set; }

}
ModelOutput
public class ModelOutput
{
    [ColumnName(@"Text")]
    public float[] Text { get; set; }

    [ColumnName(@"ContainsToken")]
    public uint ContainsToken { get; set; }

    [ColumnName(@"Features")]
    public float[] Features { get; set; }

    [ColumnName(@"PredictedLabel")]
    public string PredictedLabel { get; set; }

    [ColumnName(@"Score")]
    public float[] Score { get; set; }

}

Также добавим логику в Predict Result, так как теперь модель будет отдавать метку Clean с положительным значением.

public class PredictResult
{
    /// <summary>
    /// Name of the predictor
    /// </summary>
    public string? Name { get; set; }
    /// <summary>
    /// Score of the prediction
    /// </summary>
    public float Score { get; set; }

    /// <summary>
    /// Is credential?
    /// </summary>
    public bool IsTrue
    {
        get
        {
            if (Name != null && !Name.Contains("Clean") && Score > 0)
            {
                return true;
            }
            else if (Name != null && Name.Contains("Clean") )
            {
                return false;
            }
            return Score > 0;
        }
    }

    /// <summary>
    /// Duration of the prediction
    /// </summary>
    public int Duration { get; set; }
}

И добавим класс основной модели:

public class GeneralPredictor : MLPredictor, IPredictor
{
    private readonly Lazy<PredictionEngine<ModelInput, ModelOutput>> _predictEngineToken;
    private readonly Stopwatch _timer = new Stopwatch();

    public GeneralPredictor()
    {
        _predictEngineToken = new Lazy<PredictionEngine<ModelInput, ModelOutput>>(() => 
            CreatePredictEngine("GeneralModel.mlnet"), true);
    }
    public async Task<PredictResult> PredictAsync(string str)
    {
        _timer.Reset();
        _timer.Start();
        var predictResult = Predict(new ModelInput
        {
            Text = str
        }, _predictEngineToken);
        _timer.Stop();
        return new PredictResult
        {
            Name = $"ML {predictResult.PredictedLabel}",
            Score = predictResult.Score.Max(),
            Duration = (int)_timer.ElapsedMilliseconds
        };
    }

}

После запуска и проверки работы нужно провести несколько итераций дообучения для получения лучшего результата. Мы провели четыре итерации, пока результаты предсказания моделью перестали отличаться от регулярных выражений. Это означает, что модель достигла определенного уровня стабильности и ее дальнейшее обучение не приводит к значительному улучшению результатов.

Масштабируем

На предыдущем этапе у нас был WEB API сервер, способный анализировать входящую строку на основании регулярных выражений и машинного анализа. Но текущая реализация не потокобезопасна. Если произойдет одновременный вызов Predict, во втором и последующих потоках возникнет исключение. 

Для работы с многопоточностью необходимо подключить библиотеку Microsoft.Extensions.ML и изменить Program.cs, добавив AddPredictionEnginePool. PredictionEnginePool позволяет создавать и переиспользовать экземпляры PredictionEngine:

var builder = WebApplication.CreateBuilder(args);

builder.Services.AddEndpointsApiExplorer();
builder.Services.AddSwaggerGen();
builder.Services.AddSingleton<RegexPredictor, RegexPredictor>();
builder.Services.AddScoped<GeneralPredictor, GeneralPredictor>();
builder.Services.AddPredictionEnginePool<ModelInput, ModelOutput>()
    .FromFile(filePath: "GeneralModel.mlnet", watchForChanges: true);

var app = builder.Build();
app.UseSwagger();
app.UseSwaggerUI();
app.MapGet("/predictCred", async (RegexPredictor regexPredictor, GeneralPredictor generalPredictor, string str) =>
{
    var predictors = new List<IPredictor>
    {
        regexPredictor,
        generalPredictor
    };
    var combinePredictor = new CombinePredictor(predictors);
    var result = await combinePredictor.PredictAsync(str);
    return result;
}).WithName("Predict string");
app.Run();

Также изменим класс GeneralPredictor, добавив в конструктор PredictionEnginePool

public class GeneralPredictor : MLPredictor, IPredictor
{
    private PredictionEnginePool<ModelInput, ModelOutput> _predictionEnginePool;
    private readonly Stopwatch _timer = new Stopwatch();

    public GeneralPredictor(PredictionEnginePool<ModelInput, ModelOutput> predictionEnginePool)
    {
        _predictionEnginePool = predictionEnginePool;
    }
    public async Task<PredictResult> PredictAsync(string str)
    {
        _timer.Reset();
        _timer.Start();
        var predictResult = Predict(new ModelInput
        {
            Text = str
        }, _predictionEnginePool);
        _timer.Stop();
        return new PredictResult
        {
            Name = $"ML {predictResult.PredictedLabel}",
            Score = predictResult.Score.Max(),
            Duration = (int)_timer.ElapsedMilliseconds
        };
    }

}

И в базовом классе MLPredictor изменим метод Predict:

public abstract class MLPredictor
{
    /// <summary>
    /// Predict the input string
    /// </summary>
    /// <param name="input"></param>
    /// <param name="predictEnginePool"></param>
    /// <returns>ModelOutput</returns>
    public ModelOutput Predict(ModelInput input, PredictionEnginePool<ModelInput, ModelOutput> predictEnginePool)
    {
        return predictEnginePool.Predict(input);
    }

}

Сравниваем с регулярными выражениями

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

Регулярные выражения

ML

Скорость 

В зависимости от количества правил, в среднем 15мс при несоответствии всем регулярным выражениям.

Зависит от того, есть ли свободный PredictionEnginePool или нужно создавать новый. При создании 140 мс, при использовании текущего 20мс

Точность

Высокая точность при большом количестве правил, но не находит соответствие с опечатками

Зависит от данных для обучения. Находит соответствие с опечатками, пример:

Пороль: P@ssw0rd

Сложность изменения

Необходимо добавить новую строку в rules.yaml

Необходимо добавить данные для обучения, обучить модель

Ложноположительные значения

Пример ложноположительной строки:

Пароль: скину в личку

В зависимости от данных для обучения. 

У себя мы решили: комбинируем оба подхода, выполняя оценку параллельно и выдавая лучший результат. Регулярные выражения хорошо подходят для тех данных, которые у нас встречаются редко, например, API-ключей: за всю историю тикетов в Jira мы нашли всего пару строк. 

Планы по развитию

За полгода работы нашего сканера на реальных данных мы обнаружили множество интересных вещей. Найдены были пароли, basic- и bearer-ключи, строка подключения к PostgreSQL, несколько строк подключения к FTP, а также учетные данные, хранящиеся в переменных окружения. Это еще раз показывает, как важно внимательно относиться к безопасности при работе с чувствительной информацией и обучать сотрудников основам кибербезопасности.

Нам еще есть куда расти. В планах на ближайшее время:

  1. Подключение дополнительных источников информации — корпоративная Wiki, общие файловые шары.

  2. Добавление новых типов информации: паспортные данные, СНИЛС и так далее.

Tags:
Hubs:
Total votes 30: ↑30 and ↓0+30
Comments13

Articles

Information

Website
tech.lamoda.ru
Registered
Founded
Employees
5,001–10,000 employees
Location
Россия