[NetClassifierUpdater] use custom fields from netbox by default

init

2 files changed, 168 insertions, 68 deletions

diff --git a/ydb/core/cms/console/net_classifier_updater.cpp b/ydb/core/cms/console/net_classifier_updater.cpp
index d863c748a0..daee724800 100644
--- a/ydb/core/cms/console/net_classifier_updater.cpp
+++ b/ydb/core/cms/console/net_classifier_updater.cpp
@@ -176,7 +176,7 @@ private:
 
     auto FormNetDataFromJson(TStringBuf jsonData) const {
         NKikimrNetClassifier::TNetData netData;
-        TVector<TString> tagsToFilter(UpdaterConfig().GetNetBoxTags().begin(), UpdaterConfig().GetNetBoxTags().end());
+        THashSet<TString> tagsToFilter(UpdaterConfig().GetNetBoxTags().begin(), UpdaterConfig().GetNetBoxTags().end());
         NJson::TJsonValue value;
         bool res = NJson::ReadJsonTree(jsonData, &value);
         if (!res)
@@ -188,20 +188,31 @@ private:
                 return NKikimrNetClassifier::TNetData{};
             TString mask = v["prefix"].GetString();
 
-            if (!v["tags"].IsArray() || v["tags"].GetArray().size() == 0)
-                return NKikimrNetClassifier::TNetData{};
-            const auto& tags = v["tags"].GetArray();
             TString label;
-            for (auto& tag : tags) {
-                if (!tag.IsString())
+            auto customFields = v.GetValueByPath("custom_fields");
+            if (customFields) {
+                if (!customFields->IsMap() || !(*customFields)["owner"].IsString()) {
                     return NKikimrNetClassifier::TNetData{};
-                if (std::count(tagsToFilter.begin(), tagsToFilter.end(), tag.GetString())) {
-                    label = tag.GetString();
-                    break;
                 }
-            }
-            if (tagsToFilter.empty()) {
-                label = tags.front().GetString();
+                auto owner = (*customFields)["owner"].GetString();
+                if (tagsToFilter.empty() || tagsToFilter.contains(owner)) {
+                    label = owner;
+                }
+            } else {
+                if (!v["tags"].IsArray() || v["tags"].GetArray().size() == 0)
+                    return NKikimrNetClassifier::TNetData{};
+                const auto& tags = v["tags"].GetArray();
+                for (auto& tag : tags) {
+                    if (!tag.IsString())
+                        return NKikimrNetClassifier::TNetData{};
+                    if (tagsToFilter.contains(tag.GetString())) {
+                        label = tag.GetString();
+                        break;
+                    }
+                }
+                if (tagsToFilter.empty()) {
+                    label = tags.front().GetString();
+                }
             }
             if (!label) {
                 continue;
diff --git a/ydb/core/cms/console/net_classifier_updater_ut.cpp b/ydb/core/cms/console/net_classifier_updater_ut.cpp
index b340b1c317..7ebee659f8 100644
--- a/ydb/core/cms/console/net_classifier_updater_ut.cpp
+++ b/ydb/core/cms/console/net_classifier_updater_ut.cpp
@@ -189,68 +189,157 @@ Y_UNIT_TEST_SUITE(TNetClassifierUpdaterTest) {
         TestGetUpdatesFromHttpServer(ConvertToJson(netData), netData, TNetClassifierUpdaterConfig::NETBOX);
     }
 
-    Y_UNIT_TEST(TestFiltrationByNetboxTags) {
-    const TString netboxResponce = "{   \
-        \"count\": 5,                   \
-        \"results\": [                  \
-            {\"prefix\": \"5.45.192.0/18\", \"tags\": [\"asd\", \"zxcv\"]},     \
-            {\"prefix\": \"5.255.192.0/18\", \"tags\": [\"zxcv\", \"asd\"]},    \
-            {\"prefix\": \"37.9.64.0/18\", \"tags\": [\"zxcv\"]},               \
-            {\"prefix\": \"95.108.128.0/17\", \"tags\": [\"asd\"]},             \
-            {\"prefix\": \"172.24.0.0/13\", \"tags\": [\"qwerty\"]}            \
-        ]}";
-
+    void RunNetBoxTest(const TString& netboxResponce, const TVector<TString>& tags, const TVector<std::pair<TString, TString>>& expected) {
         auto addMask = [](NKikimrNetClassifier::TNetData& data, const TString& mask, const TString& label) {
             auto& subnet = *data.AddSubnets();
             subnet.SetMask(mask);
             subnet.SetLabel(label);
         };
-
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "5.45.192.0/18", "asd");
-            addMask(data, "5.255.192.0/18", "zxcv");
-            addMask(data, "37.9.64.0/18", "zxcv");
-            addMask(data, "95.108.128.0/17", "asd");
-            addMask(data, "172.24.0.0/13", "qwerty");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX);
-        }
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "5.45.192.0/18", "asd");
-            addMask(data, "5.255.192.0/18", "zxcv");
-            addMask(data, "37.9.64.0/18", "zxcv");
-            addMask(data, "95.108.128.0/17", "asd");
-            addMask(data, "172.24.0.0/13", "qwerty");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, {"asd", "zxcv", "qwerty", "faketag"});
-        }
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "5.45.192.0/18", "asd");
-            addMask(data, "5.255.192.0/18", "zxcv");
-            addMask(data, "37.9.64.0/18", "zxcv");
-            addMask(data, "95.108.128.0/17", "asd");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, {"zxcv", "asd"});
-        }
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "5.45.192.0/18", "zxcv");
-            addMask(data, "5.255.192.0/18", "zxcv");
-            addMask(data, "37.9.64.0/18", "zxcv");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, {"zxcv"});
-        }
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "5.45.192.0/18", "asd");
-            addMask(data, "5.255.192.0/18", "asd");
-            addMask(data, "95.108.128.0/17", "asd");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, {"asd"});
-        }
-        {
-            NKikimrNetClassifier::TNetData data;
-            addMask(data, "172.24.0.0/13", "qwerty");
-            TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, {"qwerty"});
+        NKikimrNetClassifier::TNetData data;
+        for (auto& net : expected) {
+            addMask(data, net.first, net.second);
         }
+        TestGetUpdatesFromHttpServer(netboxResponce, data, TNetClassifierUpdaterConfig::NETBOX, tags);
+    }
+    
+    void RunNetBoxCommonTests(const TString& netboxResponce) {
+        RunNetBoxTest(
+            netboxResponce,
+            {},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"},
+                {"95.108.128.0/17", "asd"},
+                {"172.24.0.0/13", "qwerty"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"asd2", "asd", "zxcv", "qwerty", "faketag"},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"},
+                {"95.108.128.0/17", "asd"},
+                {"172.24.0.0/13", "qwerty"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"asd", "zxcv"},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"},
+                {"95.108.128.0/17", "asd"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"zxcv", "asd"},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"},
+                {"95.108.128.0/17", "asd"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"qwerty"},
+            {
+                {"172.24.0.0/13", "qwerty"}
+            }
+        );
+    }
+
+    void RunTestWithCastomFields(const TString& netboxResponce) {
+        RunNetBoxCommonTests(netboxResponce);
+        RunNetBoxTest(
+            netboxResponce,
+            {"asd"},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"95.108.128.0/17", "asd"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"zxcv"},
+            {
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"}
+            }
+        );
+    }
+    
+    Y_UNIT_TEST(TestFiltrationByNetboxCustomFieldsAndTags) {
+        const TString netboxResponce = R"__(
+            {
+                "count": 5,
+                "results": [
+                    {"prefix": "5.45.192.0/18", "custom_fields": {"owner": "asd"}, "tags": ["asd", "zxcv"]},
+                    {"prefix": "5.255.192.0/18", "custom_fields": {"owner": "zxcv"}, "tags": ["zxcv", "asd"]},
+                    {"prefix": "37.9.64.0/18", "custom_fields": {"owner": "zxcv"}, "tags": ["zxcv"]},
+                    {"prefix": "95.108.128.0/17", "custom_fields": {"owner": "asd"}, "tags": ["asd"]},
+                    {"prefix": "172.24.0.0/13", "custom_fields": {"owner": "qwerty"}, "tags": ["qwerty"]}
+                ]
+            }
+        )__";
+        RunTestWithCastomFields(netboxResponce);
+    }
+    
+    Y_UNIT_TEST(TestFiltrationByNetboxCustomFieldsOnly) {
+        const TString netboxResponce = R"__(
+            {
+                "count": 5,
+                "results": [
+                    {"prefix": "5.45.192.0/18", "custom_fields": {"owner": "asd"}},
+                    {"prefix": "5.255.192.0/18", "custom_fields": {"owner": "zxcv"}},
+                    {"prefix": "37.9.64.0/18", "custom_fields": {"owner": "zxcv"}},
+                    {"prefix": "95.108.128.0/17", "custom_fields": {"owner": "asd"}},
+                    {"prefix": "172.24.0.0/13", "custom_fields": {"owner": "qwerty"}}
+                ]
+            }
+        )__";
+        
+        RunTestWithCastomFields(netboxResponce);
+    }
+
+    Y_UNIT_TEST(TestFiltrationByNetboxTags) {
+        const TString netboxResponce = R"__(
+            {
+                "count": 5,
+                "results": [
+                    {"prefix": "5.45.192.0/18", "tags": ["asd", "zxcv"]},
+                    {"prefix": "5.255.192.0/18", "tags": ["zxcv", "asd"]},
+                    {"prefix": "37.9.64.0/18", "tags": ["zxcv"]},
+                    {"prefix": "95.108.128.0/17", "tags": ["asd"]},
+                    {"prefix": "172.24.0.0/13", "tags": ["qwerty"]}
+                ]
+            }
+        )__";
+        RunNetBoxCommonTests(netboxResponce);
+
+        RunNetBoxTest(
+            netboxResponce,
+            {"asd"},
+            {
+                {"5.45.192.0/18", "asd"},
+                {"5.255.192.0/18", "asd"},
+                {"95.108.128.0/17", "asd"}
+            }
+        );
+        RunNetBoxTest(
+            netboxResponce,
+            {"zxcv"},
+            {
+                {"5.45.192.0/18", "zxcv"},
+                {"5.255.192.0/18", "zxcv"},
+                {"37.9.64.0/18", "zxcv"}
+            }
+        );
     }
 }