1from __future__ import unicode_literals
2from django import VERSION
3from ipaddress import ip_interface, ip_network
4from netaddr import EUI
5
6from django.db.models import Case, F, When
7from django.test import TestCase
8from unittest import skipIf
9
10from netfields.functions import (
11    Abbrev,
12    Broadcast,
13    Family,
14    Host,
15    Hostmask,
16    Masklen,
17    Netmask,
18    Network,
19    SetMasklen,
20    AsText,
21    IsSameFamily,
22    Merge,
23    Trunc
24)
25
26from test.models import (
27    AggregateTestChildModel,
28    AggregateTestModel,
29    CidrTestModel,
30    InetTestModel,
31    MACTestModel
32)
33
34
35class TestInetFieldFunctions(TestCase):
36    def setUp(self):
37        InetTestModel.objects.create(field='10.1.0.1/16')
38        InetTestModel.objects.create(field='2001:4f8:3:ba::1/64')
39
40    def test_abbreviate(self):
41        qs = InetTestModel.objects.annotate(abbrv=Abbrev(F('field')))
42        self.assertEqual(qs[0].abbrv, '10.1.0.1/16')
43        self.assertEqual(qs[1].abbrv, '2001:4f8:3:ba::1/64')
44
45    def test_broadcast(self):
46        qs = InetTestModel.objects.annotate(broadcast=Broadcast(F('field')))
47        self.assertEqual(qs[0].broadcast, ip_interface('10.1.255.255/16'))
48        self.assertEqual(qs[1].broadcast, ip_interface('2001:4f8:3:ba:ffff:ffff:ffff:ffff/64'))
49
50    def test_family(self):
51        qs = InetTestModel.objects.annotate(family=Family(F('field')))
52        self.assertEqual(qs[0].family, 4)
53        self.assertEqual(qs[1].family, 6)
54
55    def test_host(self):
56        qs = InetTestModel.objects.annotate(host=Host(F('field')))
57        self.assertEqual(qs[0].host, '10.1.0.1')
58        self.assertEqual(qs[1].host, '2001:4f8:3:ba::1')
59
60    def test_hostmask(self):
61        qs = InetTestModel.objects.annotate(hostmask=Hostmask(F('field')))
62        self.assertEqual(qs[0].hostmask, ip_interface('0.0.255.255'))
63        self.assertEqual(qs[1].hostmask, ip_interface('::ffff:ffff:ffff:ffff'))
64
65    def test_masklen(self):
66        qs = InetTestModel.objects.annotate(masklen=Masklen(F('field')))
67        self.assertEqual(qs[0].masklen, 16)
68        self.assertEqual(qs[1].masklen, 64)
69
70    def test_netmask(self):
71        qs = InetTestModel.objects.annotate(netmask=Netmask(F('field')))
72        self.assertEqual(qs[0].netmask, ip_interface('255.255.0.0'))
73        self.assertEqual(qs[1].netmask, ip_interface('ffff:ffff:ffff:ffff::'))
74
75    def test_network(self):
76        qs = InetTestModel.objects.annotate(network=Network(F('field')))
77        self.assertEqual(qs[0].network, ip_network('10.1.0.0/16'))
78        self.assertEqual(qs[1].network, ip_network('2001:4f8:3:ba::/64'))
79
80    def test_set_masklen(self):
81        (
82            InetTestModel.objects
83            .annotate(family=Family(F('field')))
84            .update(
85                field=Case(
86                    When(family=4, then=SetMasklen(F('field'), 24)),
87                    When(family=6, then=SetMasklen(F('field'), 120))
88                )
89            )
90        )
91        qs = InetTestModel.objects.all()
92        self.assertEqual(qs[0].field, ip_interface('10.1.0.1/24'))
93        self.assertEqual(qs[1].field, ip_interface('2001:4f8:3:ba::1/120'))
94
95    def test_as_text(self):
96        qs = InetTestModel.objects.annotate(text=AsText(F('field')))
97        self.assertEqual(qs[0].text, '10.1.0.1/16')
98        self.assertEqual(qs[1].text, '2001:4f8:3:ba::1/64')
99
100    def test_is_same_family(self):
101        parent = AggregateTestModel.objects.create(inet='0.0.0.0/0')
102        AggregateTestChildModel.objects.create(
103            parent=parent, inet='10.1.0.1/16', network='10.1.0.0/16'
104        )
105        AggregateTestChildModel.objects.create(
106            parent=parent, inet='2001:4f8:3:ba::1/64', network='2001:4f8:3:ba::/64'
107        )
108
109        qs = (
110            AggregateTestChildModel.objects.annotate(
111                is_same_family=IsSameFamily(F('inet'), F('parent__inet'))
112            )
113            .order_by('id')
114        )
115        self.assertEqual(qs[0].is_same_family, True)
116        self.assertEqual(qs[1].is_same_family, False)
117
118    def test_merge(self):
119        parent = AggregateTestModel.objects.create(inet='10.0.0.0/24')
120        AggregateTestChildModel.objects.create(
121            parent=parent, inet='10.0.1.0/24', network='10.0.0.0/23'
122        )
123
124        parent = AggregateTestModel.objects.create(inet='2001:4f8:3:ba::/64')
125        AggregateTestChildModel.objects.create(
126            parent=parent, inet='2001:4f8:3:bb::/64', network='2001:4f8:3:ba::/63'
127        )
128
129        qs = (
130            AggregateTestChildModel.objects.annotate(
131                merged=Merge(F('inet'), F('parent__inet'))
132            )
133        )
134        self.assertEqual(qs[0].merged, qs[0].network)
135        self.assertEqual(qs[1].merged, qs[1].network)
136
137
138class TestCidrFieldFunctions(TestCase):
139    def setUp(self):
140        CidrTestModel.objects.create(field='10.1.0.0/16')
141        CidrTestModel.objects.create(field='2001:4f8:3:ba::/64')
142
143    def test_abbreviate(self):
144        qs = CidrTestModel.objects.annotate(abbrv=Abbrev(F('field')))
145        self.assertEqual(qs[0].abbrv, '10.1/16')
146        self.assertEqual(qs[1].abbrv, '2001:4f8:3:ba/64')
147
148    def test_broadcast(self):
149        qs = CidrTestModel.objects.annotate(broadcast=Broadcast(F('field')))
150        self.assertEqual(qs[0].broadcast, ip_interface('10.1.255.255/16'))
151        self.assertEqual(qs[1].broadcast, ip_interface('2001:4f8:3:ba:ffff:ffff:ffff:ffff/64'))
152
153    def test_family(self):
154        qs = CidrTestModel.objects.annotate(family=Family(F('field')))
155        self.assertEqual(qs[0].family, 4)
156        self.assertEqual(qs[1].family, 6)
157
158    def test_host(self):
159        qs = CidrTestModel.objects.annotate(host=Host(F('field')))
160        self.assertEqual(qs[0].host, '10.1.0.0')
161        self.assertEqual(qs[1].host, '2001:4f8:3:ba::')
162
163    def test_hostmask(self):
164        qs = CidrTestModel.objects.annotate(hostmask=Hostmask(F('field')))
165        self.assertEqual(qs[0].hostmask, ip_interface('0.0.255.255'))
166        self.assertEqual(qs[1].hostmask, ip_interface('::ffff:ffff:ffff:ffff'))
167
168    def test_masklen(self):
169        qs = CidrTestModel.objects.annotate(masklen=Masklen(F('field')))
170        self.assertEqual(qs[0].masklen, 16)
171        self.assertEqual(qs[1].masklen, 64)
172
173    def test_netmask(self):
174        qs = CidrTestModel.objects.annotate(netmask=Netmask(F('field')))
175        self.assertEqual(qs[0].netmask, ip_interface('255.255.0.0'))
176        self.assertEqual(qs[1].netmask, ip_interface('ffff:ffff:ffff:ffff::'))
177
178    def test_network(self):
179        qs = CidrTestModel.objects.annotate(network=Network(F('field')))
180        self.assertEqual(qs[0].network, ip_network('10.1.0.0/16'))
181        self.assertEqual(qs[1].network, ip_network('2001:4f8:3:ba::/64'))
182
183    def test_set_masklen(self):
184        (
185            CidrTestModel.objects
186            .annotate(family=Family(F('field')))
187            .update(
188                field=Case(
189                    When(family=4, then=SetMasklen(F('field'), 24)),
190                    When(family=6, then=SetMasklen(F('field'), 120))
191                )
192            )
193        )
194        qs = CidrTestModel.objects.all()
195        self.assertEqual(qs[0].field, ip_network('10.1.0.0/24'))
196        self.assertEqual(qs[1].field, ip_network('2001:4f8:3:ba::/120'))
197
198    def test_as_text(self):
199        qs = CidrTestModel.objects.annotate(text=AsText(F('field')))
200        self.assertEqual(qs[0].text, '10.1.0.0/16')
201        self.assertEqual(qs[1].text, '2001:4f8:3:ba::/64')
202
203    def test_is_same_family(self):
204        parent = AggregateTestModel.objects.create(network='0.0.0.0/0')
205        AggregateTestChildModel.objects.create(
206            parent=parent, inet= '10.1.0.1/16', network='10.1.0.0/16'
207        )
208        AggregateTestChildModel.objects.create(
209            parent=parent, inet='2001:4f8:3:ba::1/64', network='2001:4f8:3:ba::/64'
210        )
211
212        qs = (
213            AggregateTestChildModel.objects.annotate(
214                is_same_family=IsSameFamily(F('network'), F('parent__network'))
215            )
216            .order_by('id')
217        )
218        self.assertEqual(qs[0].is_same_family, True)
219        self.assertEqual(qs[1].is_same_family, False)
220
221    def test_merge(self):
222        parent = AggregateTestModel.objects.create(network='10.0.0.0/24')
223        AggregateTestChildModel.objects.create(
224            parent=parent, inet='10.0.1.0/24', network='10.0.0.0/23'
225        )
226
227        parent = AggregateTestModel.objects.create(network='2001:4f8:3:ba::/64')
228        AggregateTestChildModel.objects.create(
229            parent=parent, inet='2001:4f8:3:bb::/64', network='2001:4f8:3:ba::/63'
230        )
231
232        qs = (
233            AggregateTestChildModel.objects.annotate(
234                merged=Merge(F('network'), F('parent__network'))
235            )
236        )
237        self.assertEqual(qs[0].merged, qs[0].network)
238        self.assertEqual(qs[1].merged, qs[1].network)
239
240    @skipIf(VERSION < (2, 0), 'Django unable to resolve type of num_ips to be IntegerField until 2.0.')
241    def test_read_me_example(self):
242        qs = (
243            CidrTestModel.objects.annotate(
244                family=Family(F('field')),
245                num_ips=2 ** (32 - Masklen(F('field'))),
246            )
247            .filter(family=4)
248        )
249        self.assertEqual(qs[0].num_ips, 65536)
250
251
252class TestMacFieldFunctions(TestCase):
253    def setUp(self):
254        MACTestModel.objects.create(field='aa:bb:cc:dd:ee:ff')
255
256    def test_trunc(self):
257        qs = MACTestModel.objects.annotate(trunc=Trunc(F('field')))
258        self.assertEqual(qs[0].trunc, EUI('aa:bb:cc:00:00:00'))
259