Refactoring the way classes are loaded into the CRUD APIEndpoint base class

jordan-wright · jordan-wright · commit d985b381beac · 2017-01-17T23:06:48.000-06:00
diff --git a/gophish/api/api.py b/gophish/api/api.py
@@ -11,17 +11,19 @@ class APIEndpoint:
     Represents an API endpoint for Gophish, containing common patterns
     for CRUD operations.
     """
-    def __init__(self, api, endpoint=None):
+    def __init__(self, api, endpoint=None, cls=None):
         """ Creates an instance of the APIEndpoint class.
 
         Args:
             api - Gophish.client - The authenticated REST client
             endpoint - str - The URL path to the resource endpoint
+            cls - gophish.models.Model - The Class to use when parsing results
         """
         self.api = api
         self.endpoint = endpoint
+        self._cls = cls
 
-    def get(self, cls, resource_id=None):
+    def get(self, resource_id=None):
         """ Gets the details for one or more resources by ID
         
         Args:
@@ -42,29 +44,28 @@ def get(self, cls, resource_id=None):
             return Error.parse(response.json())
 
         if resource_id:
-            return cls.parse(response.json())
+            return self._cls.parse(response.json())
 
-        return [cls.parse(resource) for resource in response.json()]
+        return [self._cls.parse(resource) for resource in response.json()]
 
-    def post(self, cls, resource):
+    def post(self, resource):
         """ Creates a new instance of the resource.
 
         Args:
-            cls - gophish.models.Model - The resource class
+            resource - gophish.models.Model - The resource instance
 
         """
         response = self.api.execute("POST", self.endpoint, json=json.dumps(resource))
         
         if not response.ok:
             return Error.parse(response.json())
 
-        return cls.parse(response.json())
+        return self._cls.parse(response.json())
 
-    def put(self, cls, resource):
+    def put(self, resource):
         """ Edits an existing resource
 
         Args:
-            cls - gophish.models.Model - The resource class
             resource - gophish.models.Model - The resource instance
         """
         
@@ -78,13 +79,12 @@ def put(self, cls, resource):
         if not respose.ok:
             return Error.parse(response.json())
 
-        return cls.parse(response.json())
+        return self._cls.parse(response.json())
 
-    def delete(self, cls, resource_id):
+    def delete(self, resource_id):
         """ Deletes an existing resource
 
         Args:
-            cls - gophish.models.Model - The resource class
             resource_id - int - The resource ID to be deleted
         """
 
@@ -95,4 +95,4 @@ def delete(self, cls, resource_id):
         if not response.ok:
             return Error.parse(response.json())
 
-        return cls.parse(response.json())
+        return self._cls.parse(response.json())
diff --git a/gophish/api/campaigns.py b/gophish/api/campaigns.py
@@ -8,7 +8,28 @@
 class API(APIEndpoint):
     def __init__(self, api, endpoint='/campaigns'):
         """ Creates a new instance of the campaigns API """
-        super(API, self).__init__(api, endpoint=endpoint)
+
+        super(API, self).__init__(api, endpoint=endpoint, cls=Campaign)
+
+    def get(self, campaign_id=None):
+        """ Gets the details for one or more campaigns by ID """
+
+        return super(API, self).get(resource_id=campaign_id)
+
+    def post(self, campaign):
+        """ Creates a new campaign """
+
+        return super(API, self).post(campaign)
+
+    def puts(self, campaign):
+        """ Edits an existing campaign """
+
+        return super(API, self).put(campaign)
+
+    def delete(self, campaign_id):
+        """ Deletes an existing campaign """
+
+        return super(API, self).delete(campaign_id)
 
     def summary(campaign_id=None):
         """ Returns the summary of one or more campaigns. """
diff --git a/gophish/api/groups.py b/gophish/api/groups.py
@@ -1,20 +1,21 @@
-class API(object):
+from gophish.models import Group
+
+class API(APIEndpoint):
     def __init__(self, api, endpoint='/groups'):
-        self.api = api
-        self.endpoint = endpoint
+        super(API, self).__init__(api, endpoint=endpoint, cls=Group)
 
     def get(group_id=None):
         """ Gets one or more groups """
-        raise NotImplementedError
+        return super(API, self).get(resource_id=group_id)
 
     def post(group):
         """ Creates a new group """
-        raise NotImplementedError
+        return super(API, self).post(group)
 
     def put(group):
         """ Edits a group """
-        raise NotImplementedError
+        return super(API, self).put(group)
 
     def delete(group_id):
         """ Deletes a group by ID """
-        raise NotImplementedError
+        return super(API, self).put(group_id)
diff --git a/gophish/api/pages.py b/gophish/api/pages.py
@@ -1,20 +1,26 @@
-class API(object):
+from gophish.models import Page
+from gophish.api import APIEndpoint
+
+class API(APIEndpoint):
     def __init__(self, api, endpoint='/pages'):
-        self.api = api
-        self.endpoint = endpoint
+        super(API, self).__init__(api, endpoint=endpoint, cls=Page)
 
-    def get(page_id=None):
+    def get(self, page_id=None):
         """ Gets one or more pages """
-        raise NotImplementedError
+        
+        return self.super(API, self).get(resource_id=page_id)
 
-    def post(page):
+    def post(self, page):
         """ Creates a new page """
-        raise NotImplementedError
 
-    def put(page):
+        return super(API, self).post(page)
+
+    def put(self, page):
         """ Edits a page """
-        raise NotImplementedError
 
-    def delete(page_id):
+        return super(API, self).put(put)
+
+    def delete(self, page_id):
         """ Deletes a page by ID """
-        raise NotImplementedError
+
+        return super(API, self).delete(page_id)
diff --git a/gophish/api/smtp.py b/gophish/api/smtp.py
@@ -1,20 +1,26 @@
-class API(object):
+from gophish.api import APIEndpoint
+from gophish.models import SMTP
+
+class API(APIEndpoint):
     def __init__(self, api, endpoint='/smtp'):
-        self.api = api
-        self.endpoint = endpoint
+        super(API, self).__init__(api, endpoint=endpoint, cls=SMTP)
 
-    def get(smtp_id=None):
+    def get(self, smtp_id=None):
         """ Gets one or more SMTP sending profiles """
-        raise NotImplementedError
 
-    def post(smtp):
+        return super(API, self).get(resource_id=smtp_id)
+
+    def post(self, smtp):
         """ Creates a new SMTP sending profile """
-        raise NotImplementedError
 
-    def put(smtp):
+        return super(API, self).post(smtp)
+
+    def put(self, smtp):
         """ Edits an SMTP sending profile """
-        raise NotImplementedError
 
-    def delete(smtp_id):
+        return super(API, self).put(smtp)
+
+    def delete(self, smtp_id):
         """ Deletes an SMTP sending profile by ID """
-        raise NotImplementedError
+
+        return super(API, self).delete(smtp_id)
diff --git a/gophish/api/templates.py b/gophish/api/templates.py
@@ -1,20 +1,26 @@
-class API(object):
+from gophish.models import Template
+from gophish.api import APIEndpoint
+
+class API(APIEndpoint):
     def __init__(self, api, endpoint='/templates'):
-        self.api = api
-        self.endpoint = endpoint
+        super(API, self).__init__(api, endpoint=endpoint, cls=Template)
 
-    def get(template_id=None):
+    def get(self, template_id=None):
         """ Gets one or more templates """
-        raise NotImplementedError
+        
+        return super(API, self).get(resource_id=campaign_id)
 
-    def post(template):
+    def post(self, template):
         """ Creates a new template """
-        raise NotImplementedError
 
-    def put(template):
+        return super(API, self).post(template)
+
+    def put(self, template):
         """ Edits a template """
-        raise NotImplementedError
 
-    def delete(template_id):
+        return super(API, self).put(template)
+
+    def delete(self, template_id):
         """ Deletes a template by ID """
-        raise NotImplementedError
+
+        return super(API, self).delete(template_id)
