From 9990c391bc73677c50cb154c15e9d009ef2435a1 Mon Sep 17 00:00:00 2001
From: Kyle Parrish <github@arnydo.com>
Date: Fri, 18 Oct 2019 16:28:05 -0400
Subject: [PATCH] Update UmbrellaBlacklister to include FQDN and URL
 data_types.

---
 .../UmbrellaBlacklister.py                    | 31 ++++++++++++++-----
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/responders/UmbrellaBlacklister/UmbrellaBlacklister.py b/responders/UmbrellaBlacklister/UmbrellaBlacklister.py
index 31a4cd1..2ac802c 100644
--- a/responders/UmbrellaBlacklister/UmbrellaBlacklister.py
+++ b/responders/UmbrellaBlacklister/UmbrellaBlacklister.py
@@ -5,19 +5,32 @@ from cortexutils.responder import Responder
 import requests
 from datetime import datetime
 
+
 class UmbrellaBlacklister(Responder):
     def __init__(self):
         Responder.__init__(self)
-        self.integration_url = self.get_param('config.integration_url', None, "Integration URL Missing")
+        self.integration_url = self.get_param(
+            'config.integration_url', None, "Integration URL Missing")
 
     def run(self):
         Responder.run(self)
 
-        if self.get_param('data.dataType') == 'domain':
+        data_type = self.get_param('data.dataType')
+        ioc_types = {"domain": "domain", "url": "url","fqdn": "fqdn"}
+        if data_type in ioc_types:
 
-            domain = self.get_param('data.data', None, 'No artifacts available')
+            if data_type == "domain" or data_type == "fqdn":
+                domain = self.get_param(
+                    'data.data', None, 'No artifacts available')
+
+                dstUrl = "http://" + domain
+
+            elif data_type == "url":
+                dstUrl = self.get_param(
+                    'data.data', None, 'No artifacts available')
+
+                domain = dstUrl.split('/')[2]
 
-            dstUrl = "http://" + domain
             date = datetime.now().strftime("%Y-%m-%dT%XZ")
 
             headers = {
@@ -36,16 +49,18 @@ class UmbrellaBlacklister(Responder):
                 "providerName": "Security Platform"
             }
 
-            r = requests.post(self.integration_url, json=payload, headers=headers)
+            r = requests.post(self.integration_url,
+                              json=payload, headers=headers)
             if r.status_code == 200 | 202:
                 self.report({'message': 'Blacklisted in Umbrella.'})
             else:
                 self.error('Failed to add to blacklist.')
-	else:
-	    self.error('Incorrect dataType. "Domain" expexted.')
+        else:
+            self.error('Incorrect dataType. "Domain", "FQDN", or "URL" expected.')
 
     def operations(self, raw):
         return [self.build_operation('AddTagToArtifact', tag='Umbrella:blocked')]
 
+
 if __name__ == '__main__':
-        UmbrellaBlacklister().run()
+    UmbrellaBlacklister().run()