public FileClassificationResponseItem Predict(FileClassificationResponseItem file)
{
if (!File.Exists(Common.Constants.MODEL_PATH))
{
file.ErrorMessage = $"Model not found ({Common.Constants.MODEL_PATH}) - please train the model first";
return(file);
}
ITransformer mlModel;
using (var stream = new FileStream(Common.Constants.MODEL_PATH, FileMode.Open, FileAccess.Read, FileShare.Read))
{
mlModel = MlContext.Model.Load(stream, out _);
}
var predictionEngine = MlContext.Model.CreatePredictionEngine <FileData, FileDataPrediction>(mlModel);
var prediction = predictionEngine.Predict(file.ToFileData());
file.Confidence = prediction.Probability;
file.IsMalicious = prediction.PredictedLabel;
return(file);
}
private void ExtractFolder(string folderPath, string outputFile)
{
if (!Directory.Exists(folderPath))
{
Console.WriteLine($"{folderPath} does not exist");
return;
}
var files = Directory.GetFiles(folderPath);
using (var streamWriter =
new StreamWriter(Path.Combine(AppContext.BaseDirectory, $"../../../../{outputFile}")))
{
foreach (var file in files)
{
var extractedData = new FileClassificationResponseItem(File.ReadAllBytes(file)).ToFileData();
extractedData.Label = !file.Contains("clean");
streamWriter.WriteLine(extractedData.ToString());
}
}
Console.WriteLine($"Extracted {files.Length} to {outputFile}");
}
public static FileData ToFileData(this FileClassificationResponseItem fileClassification)
{
return(new FileData
{
Is64Bit = fileClassification.Is64Bit,
IsSigned = fileClassification.IsSigned,
NumberImports = fileClassification.NumImports,
NumberImportFunctions = fileClassification.NumImportFunctions,
NumberExportFunctions = fileClassification.NumExportFunctions,
FileSize = fileClassification.FileSize,
Strings = fileClassification.Strings
});
}
public FileClassificationResponseItem Post(IFormFile file)
{
if (file == null)
{
return(null);
}
var fileBytes = GetBytesFromPost(file);
var responseItem = new FileClassificationResponseItem(fileBytes);
return(_predictor.Predict(responseItem));
}
