Merge pull request #87 from elghali97/main

Enhancement - AWS - Implement dynamic way to load IP of the metastore
databricks · Sep 9, 2024 · 8f5356c · 8f5356c
2 parents ed54558 + 290b0fc
commit 8f5356c
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 68 deletions.
diff --git a/aws-gov/tf/modules/sra/data_plane_hardening/firewall/firewall.tf b/aws-gov/tf/modules/sra/data_plane_hardening/firewall/firewall.tf
@@ -186,13 +186,8 @@ resource "aws_networkfirewall_rule_group" "databricks_fqdn_allowlist" {
   }
 }
 
-// Data for IP allow list
-data "external" "metastore_ip" {
-  program = ["sh", "${path.module}/metastore_ip.sh"]
-
-  query = {
-    metastore_domain = var.hive_metastore_fqdn
-  }
+data "dns_a_record_set" "metastore_dns" {
+  host     = var.hive_metastore_fqdn
 }
 
 // JDBC Firewall group IP allow list
@@ -205,19 +200,22 @@ resource "aws_networkfirewall_rule_group" "databricks_metastore_allowlist" {
       rule_order = "STRICT_ORDER"
     }
     rules_source {
-      stateful_rule {
-        action = "PASS"
-        header {
-          destination      = data.external.metastore_ip.result["ip"]
-          destination_port = 3306
-          direction        = "FORWARD"
-          protocol         = "TCP"
-          source           = "ANY"
-          source_port      = "ANY"
-        }
-        rule_option {
-          keyword  = "sid"
-          settings = ["1"]
+      dynamic "stateful_rule" {
+        for_each = toset(data.dns_a_record_set.metastore_dns.addrs)
+        content {
+          action = "PASS"
+          header {
+              destination      = stateful_rule.value
+            destination_port = 3306
+            direction        = "FORWARD"
+            protocol         = "TCP"
+            source           = "ANY"
+            source_port      = "ANY"
+          }
+          rule_option {
+            keyword  = "sid"
+            settings = ["1"]
+          }
         }
       }
       stateful_rule {
@@ -265,7 +263,6 @@ resource "aws_networkfirewall_firewall_policy" "databricks_nfw_policy" {
       priority     = 2
       resource_arn = aws_networkfirewall_rule_group.databricks_metastore_allowlist.arn
     }
-
   }
 
   tags = {

diff --git a/aws-gov/tf/modules/sra/data_plane_hardening/firewall/metastore_ip.sh b/aws-gov/tf/modules/sra/data_plane_hardening/firewall/metastore_ip.sh
diff --git a/aws-gov/tf/modules/sra/data_plane_hardening/firewall/provider.tf b/aws-gov/tf/modules/sra/data_plane_hardening/firewall/provider.tf
@@ -3,5 +3,8 @@ terraform {
     aws = {
       source = "hashicorp/aws"
     }
+    dns = {
+      source = "hashicorp/dns"
+    }
   }
 }
diff --git a/aws/tf/modules/sra/data_plane_hardening/firewall/firewall.tf b/aws/tf/modules/sra/data_plane_hardening/firewall/firewall.tf
@@ -186,13 +186,8 @@ resource "aws_networkfirewall_rule_group" "databricks_fqdn_allowlist" {
   }
 }
 
-// Data for IP allow list
-data "external" "metastore_ip" {
-  program = ["sh", "${path.module}/metastore_ip.sh"]
-
-  query = {
-    metastore_domain = var.hive_metastore_fqdn
-  }
+data "dns_a_record_set" "metastore_dns" {
+  host     = var.hive_metastore_fqdn
 }
 
 // JDBC Firewall group IP allow list
@@ -205,19 +200,22 @@ resource "aws_networkfirewall_rule_group" "databricks_metastore_allowlist" {
       rule_order = "STRICT_ORDER"
     }
     rules_source {
-      stateful_rule {
-        action = "PASS"
-        header {
-          destination      = data.external.metastore_ip.result["ip"]
-          destination_port = 3306
-          direction        = "FORWARD"
-          protocol         = "TCP"
-          source           = "ANY"
-          source_port      = "ANY"
-        }
-        rule_option {
-          keyword  = "sid"
-          settings = ["1"]
+      dynamic "stateful_rule" {
+        for_each = toset(data.dns_a_record_set.metastore_dns.addrs)
+        content {
+          action = "PASS"
+          header {
+            destination      = stateful_rule.value
+            destination_port = 3306
+            direction        = "FORWARD"
+            protocol         = "TCP"
+            source           = "ANY"
+            source_port      = "ANY"
+          }
+          rule_option {
+            keyword  = "sid"
+            settings = ["1"]
+          }
         }
       }
       stateful_rule {
@@ -288,4 +286,4 @@ resource "aws_networkfirewall_firewall" "nfw" {
     Name    = "${var.resource_prefix}-${var.region}-databricks-nfw"
     Project = var.resource_prefix
   }
-}
+}
diff --git a/aws/tf/modules/sra/data_plane_hardening/firewall/metastore_ip.sh b/aws/tf/modules/sra/data_plane_hardening/firewall/metastore_ip.sh
diff --git a/aws/tf/modules/sra/data_plane_hardening/firewall/provider.tf b/aws/tf/modules/sra/data_plane_hardening/firewall/provider.tf
@@ -3,5 +3,8 @@ terraform {
     aws = {
       source = "hashicorp/aws"
     }
+    dns = {
+      source = "hashicorp/dns"
+    }
   }
 }