[FEATURE]ip_plugin support CIDR addr_format

4 files changed, 67 insertions, 16 deletions

diff --git a/src/maat_ip_plugin.c b/src/maat_ip_plugin.c
index 8b82da6..c2055dd 100644
--- a/src/maat_ip_plugin.c
+++ b/src/maat_ip_plugin.c
@@ -26,6 +26,7 @@ struct ip_plugin_schema {
 	int ip_type_column;
 	int start_ip_column;
 	int end_ip_column;
+    int addr_format_column;
 	int rule_tag_column;
     int gc_timeout_s;
     int table_id;
@@ -117,6 +118,11 @@ void *ip_plugin_schema_new(cJSON *json, struct table_manager *tbl_mgr,
         goto error;
     }
 
+    custom_item = cJSON_GetObjectItem(item, "addr_format");
+    if (custom_item != NULL && custom_item->type == cJSON_Number) {
+        schema->addr_format_column = custom_item->valueint;
+    }
+
     // rule_tag is optional
     custom_item = cJSON_GetObjectItem(item, "rule_tag");
     if (custom_item != NULL && custom_item->type == cJSON_Number) {
@@ -202,6 +208,7 @@ ip_plugin_rule_new(struct ip_plugin_schema *schema, const char *table_name,
     size_t column_len = 0;
     char start_ip_str[40] = {0};
     char end_ip_str[40] = {0};
+    char addr_format[16] = {"range"};
     struct ip_rule *ip_plugin_rule = ALLOC(struct ip_rule, 1);
 
     ret = get_column_pos(line, schema->item_id_column,
@@ -232,6 +239,28 @@ ip_plugin_rule_new(struct ip_plugin_schema *schema, const char *table_name,
         goto error;
     }
 
+    if (schema->addr_format_column > 0) {
+        ret = get_column_pos(line, schema->addr_format_column,
+                             &column_offset, &column_len);
+        if (ret < 0) {
+            log_fatal(logger, MODULE_IP_PLUGIN,
+                      "[%s:%d] ip_plugin table:<%s> has no addr_format column in line:%s",
+                      __FUNCTION__, __LINE__, table_name, line);
+            goto error;
+        }
+
+        memset(addr_format, 0, sizeof(addr_format));
+        memcpy(addr_format, (line + column_offset), column_len);
+    }
+
+    if (IP_FORMAT_UNKNOWN == ip_format_str2int(addr_format)) {
+        log_fatal(logger, MODULE_IP_PLUGIN,
+                  "[%s:%d] ip_plugin table(table_id:%d) has invalid addr_format,"
+                  " should be range/CIDR, line:%s",
+                  __FUNCTION__, __LINE__, schema->table_id, line);
+        goto error;
+    }
+
     ret = get_column_pos(line, schema->start_ip_column,
                          &column_offset, &column_len);
     if (ret < 0) {
@@ -253,7 +282,7 @@ ip_plugin_rule_new(struct ip_plugin_schema *schema, const char *table_name,
     strncpy(end_ip_str, line + column_offset, column_len);
 
     if (IPv4 == ip_plugin_rule->type) {
-        ret = ip_format2range(ip_plugin_rule->type, IP_FORMAT_RANGE,
+        ret = ip_format2range(ip_plugin_rule->type, ip_format_str2int(addr_format),
                               start_ip_str, end_ip_str, 
                               &ip_plugin_rule->ipv4_rule.start_ip, 
                               &ip_plugin_rule->ipv4_rule.end_ip);
@@ -266,7 +295,7 @@ ip_plugin_rule_new(struct ip_plugin_schema *schema, const char *table_name,
         }
     } else {
         //ipv6
-        ret = ip_format2range(ip_plugin_rule->type, IP_FORMAT_RANGE,
+        ret = ip_format2range(ip_plugin_rule->type, ip_format_str2int(addr_format),
                               start_ip_str, end_ip_str, 
                               ip_plugin_rule->ipv6_rule.start_ip,
                               ip_plugin_rule->ipv6_rule.end_ip);
diff --git a/test/maat_framework_gtest.cpp b/test/maat_framework_gtest.cpp
index 42b9c75..91f7921 100644
--- a/test/maat_framework_gtest.cpp
+++ b/test/maat_framework_gtest.cpp
@@ -8444,20 +8444,20 @@ TEST_F(MaatCmdTest, PluginEXData) {
 }
 
 TEST_F(MaatCmdTest, UpdateIPPlugin) {
-    const char *table_name = "TEST_IP_PLUGIN_WITH_EXDATA";
+    const char *table_name = "TEST_IP_PLUGIN_WITH_ADDR_FORMAT";
 	const int TEST_CMD_LINE_NUM = 4;
 	struct maat *maat_inst = MaatCmdTest::_shared_maat_inst;
     int *ex_data_counter = MaatCmdTest::_ex_data_counter;
 	const char *table_line_add[TEST_CMD_LINE_NUM] = {
-        "101\t4\t192.168.30.99\t192.168.30.101\tSomething-like-json\t1",
-		"102\t4\t192.168.30.90\t192.168.30.128\tBigger-range-should-in-the-back\t1",
-        "103\t6\t2001:db8:1234::\t2001:db8:1235::\tBigger-range-should-in-the-back\t1",
-		"104\t6\t2001:db8:1234::1\t2001:db8:1234::5210\tSomething-like-json\t1"};
+        "101\t4\t192.168.30.98\t31\tSomething-like-json\t1\tCIDR",
+		"102\t4\t192.168.30.90\t192.168.30.128\tBigger-range-should-in-the-back\t1\trange",
+        "103\t6\t2001:db8:1234::\t2001:db8:1235::\tBigger-range-should-in-the-back\t1\trange",
+		"104\t6\t2001:db8:1234::1\t2001:db8:1234::5210\tSomething-like-json\t1\trange"};
     const char *table_line_del[TEST_CMD_LINE_NUM] = {
-        "101\t4\t192.168.30.99\t192.168.30.101\tSomething-like-json\t0",
-		"102\t4\t192.168.30.90\t192.168.30.128\tBigger-range-should-in-the-back\t0",
-        "103\t6\t2001:db8:1234::\t2001:db8:1235::\tBigger-range-should-in-the-back\t0",
-		"104\t6\t2001:db8:1234::1\t2001:db8:1234::5210\tSomething-like-json\t0"};
+        "101\t4\t192.168.30.98\t31\tSomething-like-json\t0\tcidr",
+		"102\t4\t192.168.30.90\t192.168.30.128\tBigger-range-should-in-the-back\t0\trange",
+        "103\t6\t2001:db8:1234::\t2001:db8:1235::\tBigger-range-should-in-the-back\t0\trange",
+		"104\t6\t2001:db8:1234::1\t2001:db8:1234::5210\tSomething-like-json\t0\trange"};
 
     int table_id = maat_get_table_id(maat_inst, table_name);
     ASSERT_GT(table_id, 0);
@@ -8492,7 +8492,7 @@ TEST_F(MaatCmdTest, UpdateIPPlugin) {
     struct ip_addr ipv4, ipv6;
     struct ip_plugin_ud *results[ARRAY_SIZE];
     ipv4.ip_type = IPV4;
-    inet_pton(AF_INET, "192.168.30.100", &(ipv4.ipv4));
+    inet_pton(AF_INET, "192.168.30.99", &(ipv4.ipv4));
 
     memset(results, 0, sizeof(results));
     ret = maat_ip_plugin_table_get_ex_data(maat_inst, table_id, &ipv4,
diff --git a/test/table_info.conf b/test/table_info.conf
index 02543cf..1792213 100644
--- a/test/table_info.conf
+++ b/test/table_info.conf
@@ -719,5 +719,19 @@
         "table_name":"ASN_NOT_LOGIC",
         "table_type":"virtual",
         "physical_table":"AS_NUMBER"
+    },
+    {
+        "table_id":66,
+        "table_name":"TEST_IP_PLUGIN_WITH_ADDR_FORMAT",
+        "table_type":"ip_plugin",
+        "valid_column":6,
+        "custom": {
+            "gc_timeout_s": 3,
+            "item_id":1,
+            "ip_type":2,
+            "start_ip":3,
+            "end_ip":4,
+            "addr_format":7
+        }
     }
 ]
\ No newline at end of file
diff --git a/tools/maat_redis_tool.cpp b/tools/maat_redis_tool.cpp
index cf7317b..7391101 100644
--- a/tools/maat_redis_tool.cpp
+++ b/tools/maat_redis_tool.cpp
@@ -27,13 +27,13 @@ void maat_tool_print_usage(void)
 	printf("\t-p [port], redis port, 6379 as default.\n");
 	printf("\t-n [db], redis db, 0 as default.\n");
 	printf("\t-d [dir], dump rules from redis to [dir], %s as default.\n", redis_dump_dir);
-	printf("\t-k yes, try to execute the 'keys EFFECTIVE_RULE:*' transaction, and give a specific reason if an error occurs\n");
+	printf("\t-k, try to execute the 'keys EFFECTIVE_RULE:*' transaction, and give a specific reason if an error occurs\n");
     printf("\t-u [json_file], flush redis and upload all rules to redis, confirm the risk before proceeding\n");
 	printf("\t-v [version], dump specific [version] from redis, dump latest version as default.\n");
 	printf("\t-j [payload.json], add or delete rules as maat json. Must have field compile_table field, and plugin table's valid flag must be in the last column.\n");
 	printf("\t-t [timeout], timeout config after t seconds, default is 0 which means never timeout.\n");	
 	printf("example: ./maat_redis_tool -h 127.0.0.1 -p 6379 -d %s\n",redis_dump_dir);
-	printf("    ./maat_redis_tool -h 127.0.0.1 -p 6379 -k yes\n");
+	printf("    ./maat_redis_tool -h 127.0.0.1 -p 6379 -k\n");
 	printf("	./maat_redis_tool -h 127.0.0.1 -p 6379 -j payload.json -t 300\n");
 }
 
@@ -393,6 +393,8 @@ int clear_config_in_redis(redisContext *c, struct log_handle *logger)
 int exec_keys_transaction(redisContext *c, struct log_handle *logger)
 {
 	printf("Ready to execute [keys EFFECTIVE_RULE:*] transaction...\n");
+	struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);
 
     size_t append_cmd_cnt = 0;
     int ret = redisAppendCommand(c, "MULTI");
@@ -442,7 +444,13 @@ int exec_keys_transaction(redisContext *c, struct log_handle *logger)
         return -1;
     }
 
-    printf("success!!!\n");
+	redisReply *sub_reply = reply->element[1];
+	clock_gettime(CLOCK_MONOTONIC, &end);
+	long long time_elapse_s = (end.tv_sec - start.tv_sec) +
+                              (end.tv_nsec - start.tv_nsec) / 1000000000;
+    printf("success!!! consume time:%llds, rule_num:%zu\n",
+		   time_elapse_s, sub_reply->elements);
+
     return 0;
 }
 
@@ -468,7 +476,7 @@ int main(int argc, char * argv[])
 	strncpy(dump_dir, redis_dump_dir, sizeof(dump_dir));
 	struct log_handle *logger = log_handle_create(log_path, 0);
 
-	while ((oc = getopt(argc,argv,"h:p:n:d:f:j:t:u:k:")) != -1) {
+	while ((oc = getopt(argc,argv,"h:p:n:d:f:j:t:u:k")) != -1) {
 		switch (oc) {
 			case 'h':
 				strncpy(redis_ip, optarg, sizeof(redis_ip));