diff --git a/profile_translator_blocklist/Policy.py b/profile_translator_blocklist/Policy.py index f79cbcc151d89637c860561c846a905567f4929e..2d4a535f742da46f14434e3f0ec7db578d3fe267 100644 --- a/profile_translator_blocklist/Policy.py +++ b/profile_translator_blocklist/Policy.py @@ -462,6 +462,6 @@ class Policy: """ highest_protocol = list(dict.keys(self.profile_data["protocols"]))[-1] id = highest_protocol - for _, value in dict.items(self.profile_data["protocols"][highest_protocol]): - id += f"_{value}" + for key, value in dict.items(self.profile_data["protocols"][highest_protocol]): + id += f"_{key}_{value}" return id diff --git a/profile_translator_blocklist/templates/main.c.j2 b/profile_translator_blocklist/templates/main.c.j2 index 5d8a39d4c7b88651001409d87c1f7a216ef81162..0ae6a2fd04c9323541e9544368c01fb5cff5adc0 100644 --- a/profile_translator_blocklist/templates/main.c.j2 +++ b/profile_translator_blocklist/templates/main.c.j2 @@ -1,3 +1,5 @@ +{% set use_dns = domain_names|length > 0 %} + /** * @brief SIGINT handler, flush stdout and exit. * @@ -15,7 +17,11 @@ void sigint_handler(int arg) { * @param prog program name */ void usage(char* prog) { + {% if use_dns %} fprintf(stderr, "Usage: %s [-s DNS_SERVER_IP] [-p DROP_PROBA]\n", prog); + {% else %} + fprintf(stderr, "Usage: %s [-p DROP_PROBA]\n", prog); + {% endif %} } @@ -30,7 +36,9 @@ int main(int argc, char *argv[]) { // Initialize variables int ret; + {% if use_dns %} char *dns_server_ip = "8.8.8.8"; // Default DNS server: Google Quad8 + {% endif %} // Setup SIGINT handler signal(SIGINT, sigint_handler); @@ -38,7 +46,11 @@ int main(int argc, char *argv[]) { /* COMMAND LINE ARGUMENTS */ int opt; + {% if use_dns %} while ((opt = getopt(argc, argv, "hp:s:")) != -1) + {% else %} + while ((opt = getopt(argc, argv, "hp:")) != -1) + {% endif %} { switch (opt) { @@ -50,10 +62,12 @@ int main(int argc, char *argv[]) { /* Random verdict mode: drop probability (float between 0 and 1) */ DROP_PROBA = atof(optarg); break; + {% if use_dns %} case 's': /* IP address of the network gateway */ dns_server_ip = optarg; break; + {% endif %} default: usage(argv[0]); exit(EXIT_FAILURE); @@ -72,14 +86,14 @@ int main(int argc, char *argv[]) { /* GLOBAL STRUCTURES INITIALIZATION */ - {% if "dns" in custom_parsers or "mdns" in custom_parsers or domain_names|length > 0 %} + {% if "dns" in custom_parsers or "mdns" in custom_parsers or use_dns %} // Initialize variables for DNS dns_map = dns_map_create(); dns_message_t dns_response; ip_list_t ip_list; dns_entry_t *dns_entry; - {% if domain_names|length > 0 %} + {% if use_dns %} // Open socket for DNS int sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP); if (sockfd < 0) { @@ -147,7 +161,7 @@ int main(int argc, char *argv[]) { /* FREE MEMORY */ - {% if "dns" in custom_parsers or "mdns" in custom_parsers or domain_names|length > 0 %} + {% if "dns" in custom_parsers or "mdns" in custom_parsers or use_dns %} // Free DNS map dns_map_free(dns_map); {% endif %} diff --git a/profile_translator_blocklist/translator.py b/profile_translator_blocklist/translator.py index e164add8fa3e856ffbcd17140ed260086ed2f7bc..a0c61149d633a808eeff71b3b589a3982123985e 100644 --- a/profile_translator_blocklist/translator.py +++ b/profile_translator_blocklist/translator.py @@ -241,6 +241,7 @@ def write_firewall( def translate_policy( device: dict, policy_dict: dict, + nfqueue_name: str = None, nfqueue_id: int = 0, output_dir: str = os.getcwd(), rate: int = None, @@ -292,7 +293,8 @@ def translate_policy( parse_policy(policy_data_backward, global_accs, nfqueue_id + 1, rate, drop_proba, log_type, log_group) ## Output - write_firewall(device, global_accs, policy_name, output_dir, drop_proba, log_type, log_group, test) + nfqueue_name = policy_name if nfqueue_name is None else nfqueue_name + write_firewall(device, global_accs, nfqueue_name, output_dir, drop_proba, log_type, log_group, test) def translate_policies( @@ -325,7 +327,6 @@ def translate_policies( args = validate_args(output_dir, nfqueue_id, rate, drop_proba) output_dir = args["output_dir"] drop_proba = args["drop_proba"] - nfqueue_name = device.get("name", nfqueue_name) # Initialize loop variables nfq_id_inc = 10 @@ -361,11 +362,10 @@ def translate_policies( nfqueue_id += nfq_id_inc # Output + nfqueue_name = device.get("name", policy_name) if nfqueue_name is None else nfqueue_name write_firewall(device, global_accs, nfqueue_name, output_dir, drop_proba, log_type, log_group, test) - - def translate_profile( profile_path: str, nfqueue_name: str = None,