From d4ddeb5bec3ab010ec77fb6ad01013276887179c Mon Sep 17 00:00:00 2001
From: vgamayunov <74012931+vgamayunov@users.noreply.github.com>
Date: Mon, 18 Oct 2021 15:30:36 +0100
Subject: [PATCH] Added source IP filter setting for NSG rules (#541)

* added nsg_source_ip setting for NSG allow rules
---
 README.md      | 1 +
 config.json    | 1 +
 pyazhpc/arm.py | 4 ++++
 3 files changed, 6 insertions(+)

diff --git a/README.md b/README.md
index d9a4f3c3..f279d4bf 100644
--- a/README.md
+++ b/README.md
@@ -183,6 +183,7 @@ This dictionary describes the resources for the project.
 | **low_priority**           | Boolean flag to se Spot Instance (Eviction = Delete)                        |   no     |  False  |
 | **managed_identity**       | [Managed Identity property](#managed-identity-property) to use              |   no     |         |
 | **nsg_allow**              | Enabled pre-defined NSG rules `ssh`, `rdp`, `http`, `https`, `zcentral` or `grafana` (**vm only**) |   no     | ssh/rdp |
+| **nsg_source_ip**          | Set source IP filter for NSG Allow rule. Default is allow all               |   no     |  *      |
 | **ephemeral_os_disk**      | Set whether to use the ephermal disk for the operation system               |   no     |  False  |
 | **ephemeral_os_disk_placement** | Specifies the ephemeral disk placement. Possible values are: CacheDisk, ResourceDisk |   no     |  CacheDisk  |
 | **os_disk_size**           | OS Disk size in GB. This is only needed if you want to use a non default size or increase the OS disk size|   no     |         |
diff --git a/config.json b/config.json
index d73ebd2c..b72b04e5 100644
--- a/config.json
+++ b/config.json
@@ -75,6 +75,7 @@
                 "scope": "resource_group"
             },
             "nsg_allow": ["rdp", "ssh", "http", "https", "grafana"],
+            "nsg_source_ip": "Source IP filter for NSG Allow rules. Default: *",
             "ephemeral_os_disk": "Set whether to use the ephermal disk for the operation system; default os false",
             "ephemeral_os_disk_placement": "Specifies the ephemeral disk placement. Possible values are: CacheDisk (default), ResourceDisk",
             "os_disk_size": "OS Disk size in GB. This is only needed if you want to use a non default size or increase the OS disk size",
diff --git a/pyazhpc/arm.py b/pyazhpc/arm.py
index a1d1c041..da9ec164 100644
--- a/pyazhpc/arm.py
+++ b/pyazhpc/arm.py
@@ -472,6 +472,7 @@ class ArmTemplate:
         rpip = res.get("public_ip", False)
         rdns = res.get("dns_name", None)
         rnsgallow = res.get("nsg_allow", None)
+        rnsgsourceip = res.get("nsg_source_ip", None)
         rppg = res.get("proximity_placement_group", False)
         rppgname = cfg.get("proximity_placement_group_name", None)
         raz = res.get("availability_zones", None)
@@ -661,6 +662,9 @@ class ArmTemplate:
                         nsgrules = [ nsg_security_rules["rdp"] ]
                     else:
                         nsgrules = [ nsg_security_rules["ssh"] ]
+                if rnsgsourceip:
+                    for rule in nsgrules:
+                        rule["properties"]["sourceAddressPrefix"] = rnsgsourceip
 
                 self.resources.append({
                     "type": "Microsoft.Network/networkSecurityGroups",