From b7e71f9f39ba3860b722ee70ee6ca58647943c78 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 22 Jan 2020 17:48:03 -0500
Subject: [PATCH] Add tests for IP address vars

---
 netbox/extras/tests/test_scripts.py | 56 ++++++++++++++++++++++++++++-
 netbox/ipam/formfields.py           | 11 ++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/netbox/extras/tests/test_scripts.py b/netbox/extras/tests/test_scripts.py
index 26e12772f..6237d1d95 100644
--- a/netbox/extras/tests/test_scripts.py
+++ b/netbox/extras/tests/test_scripts.py
@@ -1,6 +1,6 @@
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.test import TestCase
-from netaddr import IPNetwork
+from netaddr import IPAddress, IPNetwork
 
 from dcim.models import DeviceRole
 from extras.scripts import *
@@ -186,6 +186,54 @@ class ScriptVariablesTest(TestCase):
         self.assertTrue(form.is_valid())
         self.assertEqual(form.cleaned_data['var1'], testfile)
 
+    def test_ipaddressvar(self):
+
+        class TestScript(Script):
+
+            var1 = IPAddressVar()
+
+        # Validate IP network enforcement
+        data = {'var1': '1.2.3'}
+        form = TestScript().as_form(data, None)
+        self.assertFalse(form.is_valid())
+        self.assertIn('var1', form.errors)
+
+        # Validate IP mask exclusion
+        data = {'var1': '192.0.2.0/24'}
+        form = TestScript().as_form(data, None)
+        self.assertFalse(form.is_valid())
+        self.assertIn('var1', form.errors)
+
+        # Validate valid data
+        data = {'var1': '192.0.2.1'}
+        form = TestScript().as_form(data, None)
+        self.assertTrue(form.is_valid())
+        self.assertEqual(form.cleaned_data['var1'], IPAddress(data['var1']))
+
+    def test_ipaddresswithmaskvar(self):
+
+        class TestScript(Script):
+
+            var1 = IPAddressWithMaskVar()
+
+        # Validate IP network enforcement
+        data = {'var1': '1.2.3'}
+        form = TestScript().as_form(data, None)
+        self.assertFalse(form.is_valid())
+        self.assertIn('var1', form.errors)
+
+        # Validate IP mask requirement
+        data = {'var1': '192.0.2.0'}
+        form = TestScript().as_form(data, None)
+        self.assertFalse(form.is_valid())
+        self.assertIn('var1', form.errors)
+
+        # Validate valid data
+        data = {'var1': '192.0.2.0/24'}
+        form = TestScript().as_form(data, None)
+        self.assertTrue(form.is_valid())
+        self.assertEqual(form.cleaned_data['var1'], IPNetwork(data['var1']))
+
     def test_ipnetworkvar(self):
 
         class TestScript(Script):
@@ -198,6 +246,12 @@ class ScriptVariablesTest(TestCase):
         self.assertFalse(form.is_valid())
         self.assertIn('var1', form.errors)
 
+        # Validate host IP check
+        data = {'var1': '192.0.2.1/24'}
+        form = TestScript().as_form(data, None)
+        self.assertFalse(form.is_valid())
+        self.assertIn('var1', form.errors)
+
         # Validate valid data
         data = {'var1': '192.0.2.0/24'}
         form = TestScript().as_form(data, None)
diff --git a/netbox/ipam/formfields.py b/netbox/ipam/formfields.py
index 1ab0a8ce3..e8d171d7f 100644
--- a/netbox/ipam/formfields.py
+++ b/netbox/ipam/formfields.py
@@ -1,5 +1,6 @@
 from django import forms
 from django.core.exceptions import ValidationError
+from django.core.validators import validate_ipv4_address, validate_ipv6_address
 from netaddr import IPAddress, IPNetwork, AddrFormatError
 
 
@@ -19,6 +20,16 @@ class IPAddressFormField(forms.Field):
         if isinstance(value, IPAddress):
             return value
 
+        # netaddr is a bit too liberal with what it accepts as a valid IP address. For example, '1.2.3' will become
+        # IPAddress('1.2.0.3'). Here, we employ Django's built-in IPv4 and IPv6 address validators as a sanity check.
+        try:
+            validate_ipv4_address(value)
+        except ValidationError:
+            try:
+                validate_ipv6_address(value)
+            except ValidationError:
+                raise ValidationError("Invalid IPv4/IPv6 address format: {}".format(value))
+
         try:
             return IPAddress(value)
         except ValueError: