static void Main(string[] args)
{
#if RELEASE
var options = CommandLineParser.Parse(args);
if (options == null)
{
return;
}
#else
var options = new ClassifierCommandLineOptions {
FileName = args[0], LogLevel = LogLevels.DEBUG
};
#endif
try
{
var classifier = new Classifier(options);
var result = classifier.Classify();
Console.WriteLine(result);
} catch (Exception ex)
{
Console.WriteLine(ex);
}
}
public static ClassifierCommandLineOptions Parse(string[] args)
{
ClassifierCommandLineOptions options = null;
var oVerbose = new Option(
"--verbose",
"Enable verbose output",
new Argument <bool>(defaultValue: false));
var oFile = new Option(
"--file",
"File to be scanned (Required)",
new Argument <string>());
var rootCommand = new RootCommand
{
Description = "File Classifier applies ML to all files to determine if it is benign or malicious"
};
rootCommand.AddOption(oFile);
rootCommand.AddOption(oVerbose);
rootCommand.TreatUnmatchedTokensAsErrors = true;
rootCommand.Argument.AddValidator(symbolResult =>
{
if (symbolResult.Children["--file"] is null)
{
return("Filename is required");
}
else
{
return(null);
}
});
rootCommand.Handler = CommandHandler.Create <string, bool>((fileName, verbose) =>
{
if (string.IsNullOrEmpty(fileName))
{
return;
}
options = new ClassifierCommandLineOptions
{
FileName = fileName,
LogLevel = LogLevels.DEBUG
};
});
rootCommand.InvokeAsync(args).Wait();
return(options);
}
private static void SanityCheckOptions(ClassifierCommandLineOptions option)
{
if (option is null)
{
throw new ArgumentNullException(nameof(option));
}
if (string.IsNullOrEmpty(option.FileName))
{
throw new ArgumentNullException(nameof(option.FileName));
}
if (!File.Exists(option.FileName))
{
throw new FileNotFoundException($"{option.FileName} was not found...");
}
}
public ClassifierResponseItem Predict(ClassifierResponseItem response, ClassifierCommandLineOptions options)
{
if (response == null)
{
throw new ArgumentNullException(nameof(response));
}
var assembly = typeof(BasePredictionData).GetTypeInfo().Assembly;
var resource = assembly.GetManifestResourceStream($"FileClassifier.lib.Models.{MODEL_NAME}");
var model = MlContext.Model.Load(resource, out var schema);
var predictor = MlContext.Model.CreatePredictionEngine <T, TK>(model);
var(data, _) = FeatureExtraction(response);
var result = predictor.Predict(data);
return(UpdateResponse(result, response, options));
}
protected override ClassifierResponseItem UpdateResponse(ClusterDataPrediction prediction, ClassifierResponseItem response, ClassifierCommandLineOptions options)
{
response.FileGroup = (FileGroupType)prediction.PredictedClusterId;
var distances = prediction.Distances.Select((t, x) => $"{(FileGroupType)x+1}:{t}").ToList();
Logger <ClassifierCommandLineOptions> .Debug($"Distances: {string.Join("|", distances)}", options);
response.UpdateStatus(ClassifierStatus.SUCCESS);
return(response);
}
protected override ClassifierResponseItem UpdateResponse(ClassificationDataPrediction prediction, ClassifierResponseItem response, ClassifierCommandLineOptions options)
{
response.Confidence = prediction.Score;
response.IsMalicious = prediction.Prediction;
return(response);
}
protected abstract ClassifierResponseItem UpdateResponse(TK prediction, ClassifierResponseItem response, ClassifierCommandLineOptions options);
public Classifier(ClassifierCommandLineOptions option)
{
SanityCheckOptions(option);
_options = option;
}
