WIF - Classifiers - Introduce IP prefix classifier

TG-36

Author: plnyrich

include/wif/classifiers/ipPrefixClassifier.hpp

@@ -0,0 +1,73 @@
+/**
+ * @file
+ * @author Richard Plny <[email protected]>
+ * @brief IP prefix classifier interface
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#pragma once
+
+#include "wif/classifiers/classifier.hpp"
+#include "wif/utils/ipPrefix.hpp"
+
+#include <algorithm>
+#include <memory>
+#include <vector>
+
+namespace WIF {
+
+/**
+ * @brief Classifier performing blocklist detection
+ */
+class IpPrefixClassifier : public Classifier {
+public:
+	/**
+	 * @brief Construct an IP Prefix Classifier object with empty blocklist
+	 */
+	IpPrefixClassifier() = default;
+
+	/**
+	 * @brief Construct a new IP Prefix Classifier object
+	 * @param blocklist std::vector of blocklisted IP prefixes
+	 */
+	IpPrefixClassifier(const std::vector<IpPrefix>& blocklist);
+
+	/**
+	 * @brief Getter for current IP prefix blocklist
+	 * @return const std::vector<IpPrefix>& current blocklist
+	 */
+	const std::vector<IpPrefix>& getBlocklist() const noexcept { return m_blocklist; }
+
+	/**
+	 * @brief Update blocklist
+	 * Old blocklist is destructed and replaced by new one
+	 * @param blocklist std::vector of new IP prefixes on blocklist
+	 */
+	void updateBlocklist(const std::vector<IpPrefix>& blocklist);
+
+	/**
+	 * @brief Classify single flowFeature object
+	 * ClfResult contains non-zero double value if flowFeatures contained a field with IP from
+	 * blocklist, otherwise the double value is set to zero
+	 *
+	 * @param flowFeatures flow features to classify
+	 * @return ClfResult result of the classification
+	 */
+	ClfResult classify(const FlowFeatures& flowFeatures) override;
+
+	/**
+	 * @brief Classify a burst of flow features
+	 *
+	 * @param burstOfFlowsFeatures the burst of flow features to classify
+	 * @return std::vector<ClfResult> std::vector<ClfResult> the results of the classification
+	 */
+	std::vector<ClfResult> classify(const std::vector<FlowFeatures>& burstOfFlowsFeatures) override;
+
+private:
+	bool findIpAddress(const IpAddress& ipAddress);
+
+	std::vector<IpPrefix> m_blocklist;
+};
+
+} // namespace WIF

src/wif/CMakeLists.txt

@@ -1,5 +1,6 @@
 set(LIBWIF_SOURCES
 	classifiers/classifier.cpp
+	classifiers/ipPrefixClassifier.cpp
 	classifiers/regexClassifier.cpp
 	classifiers/scikitMlClassifier.cpp
 	combinators/averageCombinator.cpp

src/wif/classifiers/ipPrefixClassifier.cpp

@@ -0,0 +1,51 @@
+/**
+ * @file
+ * @author Richard Plny <[email protected]>
+ * @brief IP prefix classifier implementation
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include "wif/classifiers/ipPrefixClassifier.hpp"
+
+namespace WIF {
+
+IpPrefixClassifier::IpPrefixClassifier(const std::vector<IpPrefix>& blocklist)
+{
+	updateBlocklist(blocklist);
+}
+
+void IpPrefixClassifier::updateBlocklist(const std::vector<IpPrefix>& blocklist)
+{
+	m_blocklist = blocklist;
+	std::sort(m_blocklist.begin(), m_blocklist.end());
+}
+
+ClfResult IpPrefixClassifier::classify(const FlowFeatures& flowFeatures)
+{
+	for (FeatureID featureID : getSourceFeatureIDs()) {
+		auto ipAddress = flowFeatures.get<IpAddress>(featureID);
+		if (findIpAddress(ipAddress)) {
+			return ClfResult(1.0);
+		}
+	}
+	return ClfResult(0.0);
+}
+
+std::vector<ClfResult>
+IpPrefixClassifier::classify(const std::vector<FlowFeatures>& burstOfFlowsFeatures)
+{
+	std::vector<ClfResult> burstResults;
+	burstResults.reserve(burstOfFlowsFeatures.size());
+	for (const auto& flowFeatures : burstOfFlowsFeatures) {
+		burstResults.emplace_back(classify(flowFeatures));
+	}
+	return burstResults;
+}
+
+bool IpPrefixClassifier::findIpAddress(const IpAddress& ipAddress)
+{
+	return std::binary_search(m_blocklist.begin(), m_blocklist.end(), ipAddress);
+}
+
+} // namespace WIF