Created Urls table and functions (#6)

docs/src/usage/plugins/db.md

@@ -117,6 +117,35 @@ Additionally, some properties can be overridden by the State, which allows you t
 >> >>> h.db.add_targets([{...}, {...}, {...}])
 >> ```
 
+## db.add_urls(results)
+
+> Add urls results to the database
+> 
+> | Arguments | Type | Description
+> | --- | --- | ---
+> | `results` | list(dict) | A list of dictionaries containing results from some scan, Hydra, etc.
+>
+>> Examples
+>> ```python3
+>> >>> results = [
+>> ...     {
+>> ...         "ip": "1.1.1.1",
+>> ...         "target": "7gh33tjf72",
+>> ...         "urls": [
+>> ...             {
+>> ...                 "url": "https://www.google.com",
+>> ...                 "screenshot_url": "https://imgur.com/2uregtu",
+>> ...             },
+>> ...             {
+>> ...                 "url": "https://www.ebay.com",
+>> ...                 "screenshot_url": "file:///tmp/948grt.png",
+>> ...             }
+>> ...         ]
+>> ...     }
+>> ... ]
+>> >>> h.db.add_urls(results)
+>> ```
+
 ## db.find_ips(ip, **kwargs)
 
 > Filters through all the ips to return ones which match a given criteria
@@ -151,7 +180,7 @@ Additionally, some properties can be overridden by the State, which allows you t
 >>   {
 >>     'ip': '1.2.3.4', 'source': 'hydra', 'target': '123hg912',
 >>       'ports': [
->>         { 'open': True, 'port': '443', 'protocol': 'tcp', 'screenshot_url': '', 'service': 'https - Wordpress', 'updated': 1654840021 },
+>>         { 'open': True, 'port': '443', 'protocol': 'tcp', 'service': 'https - Wordpress', 'updated': 1654840021 },
 >>         ...
 >>       ]
 >>   },
@@ -173,6 +202,35 @@ Additionally, some properties can be overridden by the State, which allows you t
 >> [<class 'synack.db.models.Target'>, ...]
 >> ```
 
+## db.find_urls(url=None, ip=None, **kwargs)
+
+> Filters through all the ports to return ones which match a given criteria
+>
+> | Argument | Type | Description
+> | --- | --- | ---
+> | `url` | str | Url hosting a service on the IP
+> | `ip` | str | IP Address to search for
+> | `kwargs` | kwargs | Any attribute of the Target Database Model (codename, slug, is_active, etc.)
+>
+>> Examples
+>> ```python3
+>> >>> h.db.find_ports(codename="SLEEPYPUPPY")
+>> [
+>>   {
+>>     'ip': '1.2.3.4',
+>>     'target': '123hg912',
+>>     'ports': [
+>>       {  
+>>         'url': 'https://www.google.com',
+>>         'screenshot_url': 'file:///tmp/2948geybu24.png'
+>>       },
+>>       ...
+>>     ]
+>>   },
+>>   ...
+>> ]
+>> ```
+
 ## db.get_config(name)
 
 > Returns a configuration from the Database.

docs/src/usage/plugins/hydra.md

@@ -29,7 +29,7 @@
 >>   {
 >>     'ip': '1.2.3.4', 'source': 'hydra', 'target': '123hg912',
 >>       'ports': [
->>         { 'open': True, 'port': '443', 'protocol': 'tcp', 'screenshot_url': '', 'service': 'https - Wordpress', 'updated': 1654840021 },
+>>         { 'open': True, 'port': '443', 'protocol': 'tcp', 'service': 'https - Wordpress', 'updated': 1654840021 },
 >>         ...
 >>       ]
 >>   },

src/synack/db/alembic/versions/0c1ac7be711c_added_url_table_deleted_url_from_port_.py

@@ -0,0 +1,34 @@
+"""Added Url table/Deleted url from Port table
+
+Revision ID: 0c1ac7be711c
+Revises: deb7dd07212c
+Create Date: 2022-06-15 03:14:49.327185
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '0c1ac7be711c'
+down_revision = 'deb7dd07212c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table('urls',
+                    sa.Column('id', sa.Integer, primary_key=True),
+                    sa.Column('ip', sa.Integer, sa.ForeignKey('ips.id')),
+                    sa.Column('url', sa.VARCHAR(1024), server_default=''),
+                    sa.Column('screenshot_url', sa.VARCHAR(1024), server_default=''))
+    with op.batch_alter_table('ports') as batch_op:
+        batch_op.drop_column('url')
+        batch_op.drop_column('screenshot_url')
+
+
+def downgrade():
+    op.drop_table('urls')
+    with op.batch_alter_table('ports') as batch_op:
+        batch_op.add_column(sa.Column('url', sa.VARCHAR(200), server_default=''))
+        batch_op.add_column(sa.Column('screenshot_url', sa.VARCHAR(1000), server_default=''))

src/synack/db/models/__init__.py

@@ -6,3 +6,4 @@
 from .ip import IP
 from .organization import Organization
 from .port import Port
+from .url import Url

src/synack/db/models/url.py

@@ -0,0 +1,18 @@
+"""db/model/url.py
+
+Database Model for the Url item
+"""
+
+import sqlalchemy as sa
+from sqlalchemy.orm import declarative_base
+from .ip import IP
+
+Base = declarative_base()
+
+
+class Url(Base):
+    __tablename__ = 'urls'
+    id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
+    ip = sa.Column(sa.Integer, sa.ForeignKey(IP.id))
+    url = sa.Column(sa.VARCHAR(1024), default="")
+    screenshot_url = sa.Column(sa.VARCHAR(1024), default="")

src/synack/plugins/db.py

@@ -15,6 +15,7 @@
 from synack.db.models import IP
 from synack.db.models import Organization
 from synack.db.models import Port
+from synack.db.models import Url
 
 from .base import Plugin
 
@@ -109,8 +110,6 @@ def add_ports(self, results):
                             port=port.get('port'),
                             protocol=port.get('protocol'),
                             service=port.get('service'),
-                            screenshot_url=port.get('screenshot_url'),
-                            url=port.get('url'),
                             ip=ip.id,
                             source=result.get('source'),
                             open=port.get('open'),
@@ -119,8 +118,6 @@ def add_ports(self, results):
                     else:
                         db_port = db_port.first()
                         db_port.service = port.get('service', db_port.service)
-                        db_port.screenshot_url = port.get('screenshot_url', db_port.screenshot_url)
-                        db_port.url = port.get('url', db_port.url)
                         db_port.open = port.get('open', db_port.open)
                         db_port.updated = port.get('updated', db_port.updated)
                     session.add(db_port)
@@ -156,6 +153,35 @@ def add_targets(self, targets, **kwargs):
         session.commit()
         session.close()
 
+    def add_urls(self, results, **kwargs):
+        self.add_ips(results)
+        session = self.Session()
+        q = session.query(Url)
+        ips = session.query(IP)
+        for result in results:
+            ip = ips.filter_by(ip=result.get('ip'))
+            if ip:
+                ip = ip.first()
+                for url in result.get('urls', []):
+                    filt = sa.and_(
+                        Url.url.like(url.get('url')),
+                        Url.ip.like(ip.id))
+                    db_url = q.filter(filt)
+                    if not db_url:
+                        db_url = Url(
+                            ip=ip.id,
+                            url=url.get('url'),
+                            screenshot_url=url.get('screenshot_url'),
+                        )
+                    else:
+                        db_url = db_url.first()
+                        db_url.ip = ip.id
+                        db_url.url = url.get('url')
+                        db_url.screenshot_url = url.get('screenshot_url')
+                    session.add(db_url)
+        session.commit()
+        session.close()
+
     def find_targets(self, **kwargs):
         session = self.Session()
         targets = session.query(Target).filter_by(**kwargs).all()
@@ -193,8 +219,6 @@ def find_ports(self, port=None, protocol=None, source=None, ip=None, **kwargs):
                 "source": port.source,
                 "open": port.open,
                 "updated": port.updated,
-                "url": port.url,
-                "screenshot_url": port.screenshot_url
             })
 
         ret = list()
@@ -235,6 +259,43 @@ def find_ips(self, ip=None, **kwargs):
 
         return ret
 
+    def find_urls(self, url=None, ip=None, **kwargs):
+        session = self.Session()
+        query = session.query(Url)
+        if url:
+            query = query.filter_by(url=url)
+
+        query = query.join(IP)
+        if ip:
+            query = query.filter_by(ip=ip)
+
+        query = query.join(Target)
+        if kwargs:
+            query = query.filter_by(**kwargs)
+
+        urls = query.all()
+
+        ips = dict()
+        for url in urls:
+            ips[url.ip] = ips.get(url.ip, list())
+            ips[url.ip].append({
+                'url': url.url,
+                'screenshot_url': url.screenshot_url,
+            })
+
+        ret = list()
+        for ip_id in ips.keys():
+            ip = session.query(IP).filter_by(id=ip_id).first()
+            ret.append({
+                "ip": ip.ip,
+                "target": ip.target,
+                "urls": ips[ip_id]
+            })
+
+        session.expunge_all()
+        session.close()
+        return ret
+
     def get_config(self, name=None):
         session = self.Session()
         config = session.query(Config).filter_by(id=1).first()
@@ -316,6 +377,13 @@ def ips(self):
         session.close()
         return ips
 
+    @property
+    def urls(self):
+        session = self.Session()
+        urls = session.query(Url).all()
+        session.close()
+        return urls
+
     @property
     def api_token(self):
         return self.get_config('api_token')

src/synack/plugins/hydra.py

@@ -55,7 +55,6 @@ def build_db_input(self, results):
                             ' - ' + \
                             h_src.get('product', {'parsed': 'unknown'})['parsed']
                         service = service.strip(' - ')
-                        screenshot_url = result['ports'][port][protocol][hydra_src].get('screenshot_key', '')
                         port_open = result['ports'][port][protocol][hydra_src]['open']['parsed']
                         epoch = datetime(1970, 1, 1)
                         try:
@@ -68,7 +67,6 @@ def build_db_input(self, results):
                             "port": port,
                             "protocol": protocol,
                             "service": service,
-                            "screenshot_url": screenshot_url,
                             "open": port_open,
                             "updated": updated
                         })