Merge pull request #223 from andrewsomething/firewall-improvements

Firewall improvements

Author: koalalorenzo

README.md

@@ -131,6 +131,36 @@ droplet = digitalocean.Droplet(token="secretspecialuniquesnowflake",
 droplet.create()
 ```
 
+### Creating a Firewall
+
+This example creates a firewall that only accepts inbound tcp traffic on port 80 from a specific load balancer and allows outbout tcp traffic on all ports to all addresses.
+
+```python
+from digitalocean import Firewall, InboundRule, OutboundRule, Destinations, Sources
+
+inbound_rule = InboundRule(protocol="tcp", ports="80",
+                           sources=Sources(
+                               load_balancer_uids=[
+                                   "4de7ac8b-495b-4884-9a69-1050c6793cd6"]
+                               )
+                           )
+
+outbound_rule = OutboundRule(protocol="tcp", ports="all",
+                             destinations=Destinations(
+                               addresses=[
+                                   "0.0.0.0/0",
+                                   "::/0"]
+                                 )
+                             )
+
+firewall = Firewall(token="secretspecialuniquesnowflake",
+                    name="new-firewall",
+                    inbound_rules=[inbound_rule],
+                    outbound_rules=[outbound_rule],
+                    droplet_ids=[8043964, 8043972])
+firewall.create()
+```
+
 ## Getting account requests/hour limits status:
 Each request will also include the rate limit information:
 

digitalocean/Firewall.py

@@ -1,16 +1,144 @@
 # -*- coding: utf-8 -*-
 from .baseapi import BaseAPI, POST, DELETE, PUT
+import jsonpickle
+
+
+class _targets(object):
+    """
+    An internal object that both `Sources` and `Destinations` derive from.
+
+    Not for direct use by end users.
+    """
+    def __init__(self, **kwargs):
+        self.addresses = []
+        self.droplet_ids = []
+        self.load_balancer_uids = []
+        self.tags = []
+
+        for attr in kwargs.keys():
+            setattr(self, attr, kwargs[attr])
+
+
+class Sources(_targets):
+    """
+    An object holding information about an InboundRule's sources.
+
+    Args:
+        addresses (obj:`list`): An array of strings containing the IPv4
+            addresses, IPv6 addresses, IPv4 CIDRs, and/or IPv6 CIDRs to which
+            the Firewall will allow traffic.
+        droplet_ids (obj:`list`): An array containing the IDs of the Droplets
+            to which the Firewall will allow traffic.
+        load_balancer_uids (obj:`list`): An array containing the IDs of the
+            Load Balancers to which the Firewall will allow traffic.
+        tags (obj:`list`): An array containing the names of Tags corresponding
+            to groups of Droplets to which the Firewall will allow traffic.
+    """
+    pass
+
+
+class Destinations(_targets):
+    """
+    An object holding information about an OutboundRule's destinations.
+
+    Args:
+        addresses (obj:`list`): An array of strings containing the IPv4
+            addresses, IPv6 addresses, IPv4 CIDRs, and/or IPv6 CIDRs to which
+            the Firewall will allow traffic.
+        droplet_ids (obj:`list`): An array containing the IDs of the Droplets
+            to which the Firewall will allow traffic.
+        load_balancer_uids (obj:`list`): An array containing the IDs of the
+            Load Balancers to which the Firewall will allow traffic.
+        tags (obj:`list`): An array containing the names of Tags corresponding
+            to groups of Droplets to which the Firewall will allow traffic.
+    """
+    pass
+
+
+class InboundRule(object):
+    """
+    An object holding information about a Firewall's inbound rule.
+
+    Args:
+        protocol (str): The type of traffic to be allowed. This may be one
+            of "tcp", "udp", or "icmp".
+        port (str): The ports on which traffic will be allowed specified as a
+            string containing a single port, a range (e.g. "8000-9000"), or
+            "all" to open all ports for a protocol.
+        sources (obj): A `Sources` object.
+    """
+    def __init__(self, protocol="", ports="", sources=""):
+        self.protocol = protocol
+        self.ports = ports
+
+        if isinstance(sources, Sources):
+            self.sources = sources
+        else:
+            for source in sources:
+                self.sources = Sources(**sources)
+
+
+class OutboundRule(object):
+    """
+    An object holding information about a Firewall's outbound rule.
+
+    Args:
+        protocol (str): The type of traffic to be allowed. This may be one
+            of "tcp", "udp", or "icmp".
+        port (str): The ports on which traffic will be allowed specified as a
+            string containing a single port, a range (e.g. "8000-9000"), or
+            "all" to open all ports for a protocol.
+        destinations (obj): A `Destinations` object.
+    """
+    def __init__(self, protocol="", ports="", destinations=""):
+        self.protocol = protocol
+        self.ports = ports
+
+        if isinstance(destinations, Destinations):
+            self.destinations = destinations
+        else:
+            for destination in destinations:
+                self.destinations = Destinations(**destinations)
 
 
 class Firewall(BaseAPI):
+    """
+    An object representing an DigitalOcean Firewall.
+
+    Attributes accepted at creation time:
+
+    Args:
+        name (str): The Firewall's name.
+        droplet_ids (obj:`list` of `int`): A list of Droplet IDs to be assigned
+            to the Firewall.
+        tags (obj:`list` of `str`):  A list Tag names to be assigned to the
+            Firewall.
+        inbound_rules (obj:`list`): A list of `InboundRules` objects
+        outbound_rules (obj:`list`): A list of `OutboundRules` objects
+
+    Attributes returned by API:
+        id (str): A UUID to identify and reference a Firewall.
+        status (str): A status string indicating the current state of the
+            Firewall. This can be "waiting", "succeeded", or "failed".
+        created_at (str): The time at which the Firewall was created.
+        name (str): The Firewall's name.
+        pending_changes (obj:`list`): Details exactly which Droplets are having
+            their security policies updated.
+        droplet_ids (obj:`list` of `int`): A list of Droplet IDs to be assigned
+            to the Firewall.
+        tags (obj:`list` of `str`):  A list Tag names to be assigned to the
+            Firewall.
+        inbound_rules (obj:`list`): A list of `InboundRules` objects
+        outbound_rules (obj:`list`): A list of `OutboundRules` objects
+    """
     def __init__(self, *args, **kwargs):
         self.id = None
         self.status = None
         self.created_at = None
         self.pending_changes = []
         self.name = None
-        self.inbound_rules = None
-        self.outbound_rules = None
+        self.inbound_rules = []
+        self.outbound_rules = []
         self.droplet_ids = None
         self.tags = None
 
@@ -25,13 +153,45 @@ def get_object(cls, api_token, firewall_id):
         firewall.load()
         return firewall
 
+    def _set_firewall_attributes(self, data):
+        self.id = data['firewall']['id']
+        self.name = data['firewall']['name']
+        self.status = data['firewall']['status']
+        self.created_at = data['firewall']['created_at']
+        self.pending_changes = data['firewall']['pending_changes']
+        self.droplet_ids = data['firewall']['droplet_ids']
+        self.tags = data['firewall']['tags']
+
+        in_rules = list()
+        for rule in data['firewall']['inbound_rules']:
+            in_rules.append(InboundRule(**rule))
+        self.inbound_rules = in_rules
+
+        out_rules = list()
+        for rule in data['firewall']['outbound_rules']:
+            out_rules.append(OutboundRule(**rule))
+        self.outbound_rules = out_rules
+
     def load(self):
         data = self.get_data("firewalls/%s" % self.id)
-        firewall_dict = data['firewall']
+        if data:
+            self._set_firewall_attributes(data)
+
+        return self
+
+    def create(self, *args, **kwargs):
+        inbound = jsonpickle.encode(self.inbound_rules, unpicklable=False)
+        outbound = jsonpickle.encode(self.outbound_rules, unpicklable=False)
+        params = {'name': self.name,
+                  'droplet_ids': self.droplet_ids,
+                  'inbound_rules': jsonpickle.decode(inbound),
+                  'outbound_rules': jsonpickle.decode(outbound),
+                  'tags': self.tags}
+
+        data = self.get_data('firewalls/', type=POST, params=params)
 
-        # Setting the attribute values
-        for attr in firewall_dict.keys():
-            setattr(self, attr, firewall_dict[attr])
+        if data:
+            self._set_firewall_attributes(data)
 
         return self
 

digitalocean/Manager.py

@@ -11,7 +11,7 @@
 from .Domain import Domain
 from .Droplet import Droplet
 from .FloatingIP import FloatingIP
-from .Firewall import Firewall
+from .Firewall import Firewall, InboundRule, OutboundRule
 from .Image import Image
 from .LoadBalancer import LoadBalancer
 from .LoadBalancer import StickySesions, HealthCheck, ForwardingRule
@@ -359,6 +359,14 @@ def get_all_firewalls(self):
         for jsoned in data['firewalls']:
             firewall = Firewall(**jsoned)
             firewall.token = self.token
+            in_rules = list()
+            for rule in jsoned['inbound_rules']:
+                in_rules.append(InboundRule(**rule))
+            firewall.inbound_rules = in_rules
+            out_rules = list()
+            for rule in jsoned['outbound_rules']:
+                out_rules.append(OutboundRule(**rule))
+            firewall.outbound_rules = out_rules
             firewalls.append(firewall)
         return firewalls
 

digitalocean/__init__.py

@@ -26,4 +26,4 @@
 from .LoadBalancer import StickySesions, ForwardingRule, HealthCheck
 from .Certificate import Certificate
 from .Snapshot import Snapshot
-from .Firewall import Firewall
+from .Firewall import Firewall, InboundRule, OutboundRule, Destinations, Sources

digitalocean/tests/test_firewall.py

@@ -41,7 +41,19 @@ def test_load(self):
         self.assertEqual(f.id, 12345)
         self.assertEqual(f.name, "firewall")
         self.assertEqual(f.status, "succeeded")
-        # TODO: Assert the inbound and outbound rules
+        self.assertEqual(f.inbound_rules[0].ports, "80")
+        self.assertEqual(f.inbound_rules[0].protocol, "tcp")
+        self.assertEqual(f.inbound_rules[0].sources.load_balancer_uids,
+                         ["12345"])
+        self.assertEqual(f.inbound_rules[0].sources.addresses, [])
+        self.assertEqual(f.inbound_rules[0].sources.tags, [])
+        self.assertEqual(f.outbound_rules[0].ports, "80")
+        self.assertEqual(f.outbound_rules[0].protocol, "tcp")
+        self.assertEqual(
+            f.outbound_rules[0].destinations.load_balancer_uids, [])
+        self.assertEqual(f.outbound_rules[0].destinations.addresses,
+                         ["0.0.0.0/0", "::/0"])
+        self.assertEqual(f.outbound_rules[0].destinations.tags, [])
         self.assertEqual(f.created_at, "2017-05-23T21:24:00Z")
         self.assertEqual(f.droplet_ids, [12345])
         self.assertEqual(f.tags, [])

digitalocean/tests/test_manager.py

@@ -562,5 +562,42 @@ def test_get_volume_snapshots(self):
         self.assertEqual(volume_snapshots[0].resource_type, 'volume')
         self.assertEqual(len(volume_snapshots[0].regions), 1)
 
+    @responses.activate
+    def test_get_firewalls(self):
+        data = self.load_from_file('firewalls/all.json')
+
+        url = self.base_url + "firewalls"
+        responses.add(responses.GET,
+                      url,
+                      body=data,
+                      status=200,
+                      content_type='application/json')
+
+        firewalls = self.manager.get_all_firewalls()
+        f = firewalls[0]
+
+        self.assert_get_url_equal(responses.calls[0].request.url, url)
+        self.assertEqual(f.id, "12345")
+        self.assertEqual(f.name, "firewall")
+        self.assertEqual(f.status, "succeeded")
+        self.assertEqual(f.inbound_rules[0].ports, "80")
+        self.assertEqual(f.inbound_rules[0].protocol, "tcp")
+        self.assertEqual(f.inbound_rules[0].sources.load_balancer_uids,
+                         ["12345"])
+        self.assertEqual(f.inbound_rules[0].sources.addresses, [])
+        self.assertEqual(f.inbound_rules[0].sources.tags, [])
+        self.assertEqual(f.outbound_rules[0].ports, "80")
+        self.assertEqual(f.outbound_rules[0].protocol, "tcp")
+        self.assertEqual(
+            f.outbound_rules[0].destinations.load_balancer_uids, [])
+        self.assertEqual(f.outbound_rules[0].destinations.addresses,
+                         ["0.0.0.0/0", "::/0"])
+        self.assertEqual(f.outbound_rules[0].destinations.tags, [])
+        self.assertEqual(f.created_at, "2017-05-23T21:24:00Z")
+        self.assertEqual(f.droplet_ids, [12345])
+        self.assertEqual(f.tags, [])
+        self.assertEqual(f.pending_changes, [])
+
+
 if __name__ == '__main__':
     unittest.main()

requirements.txt

@@ -1,4 +1,5 @@
 requests>=2.2.1
+jsonpickle
 
 # Testing requirements
 pytest