From a09faac9a76ca1df6f49cdfe2dcaa8476c093e81 Mon Sep 17 00:00:00 2001
From: kevgliss <kevgliss@gmail.com>
Date: Thu, 15 Dec 2016 10:26:59 -0800
Subject: [PATCH] Endpoint sync fixes (#604)

---
 lemur/endpoints/models.py  |  9 ++++++---
 lemur/endpoints/service.py |  3 +++
 lemur/sources/cli.py       | 14 +++++++++++++-
 lemur/sources/service.py   | 31 ++++++++-----------------------
 4 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/lemur/endpoints/models.py b/lemur/endpoints/models.py
index b4974b21..75dfb312 100644
--- a/lemur/endpoints/models.py
+++ b/lemur/endpoints/models.py
@@ -6,11 +6,14 @@
     :license: Apache, see LICENSE for more details.
 .. moduleauthor:: Kevin Glisson <kglisson@netflix.com>
 """
+import arrow
 from sqlalchemy.orm import relationship
-from sqlalchemy import Column, Integer, String, func, DateTime, PassiveDefault, Boolean, ForeignKey
+from sqlalchemy import Column, Integer, String, Boolean, ForeignKey
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.sql.expression import case
 
+from sqlalchemy_utils import ArrowType
+
 from lemur.database import db
 
 from lemur.models import policies_ciphers
@@ -64,8 +67,8 @@ class Endpoint(db.Model):
     source_id = Column(Integer, ForeignKey('sources.id'))
     sensitive = Column(Boolean, default=False)
     source = relationship('Source', back_populates='endpoints')
-    last_updated = Column(DateTime, PassiveDefault(func.now()), onupdate=func.now(), nullable=False)
-    date_created = Column(DateTime, PassiveDefault(func.now()), nullable=False)
+    last_updated = Column(ArrowType, default=arrow.utcnow, nullable=False)
+    date_created = Column(ArrowType, default=arrow.utcnow, onupdate=arrow.utcnow, nullable=False)
 
     @property
     def issues(self):
diff --git a/lemur/endpoints/service.py b/lemur/endpoints/service.py
index 1d95570c..44be9134 100644
--- a/lemur/endpoints/service.py
+++ b/lemur/endpoints/service.py
@@ -8,6 +8,8 @@
 .. moduleauthor:: Kevin Glisson <kglisson@netflix.com>
 
 """
+import arrow
+
 from flask import current_app
 
 from sqlalchemy import func
@@ -95,6 +97,7 @@ def update(endpoint_id, **kwargs):
     endpoint.policy = kwargs['policy']
     endpoint.certificate = kwargs['certificate']
     endpoint.source = kwargs['source']
+    endpoint.last_updated = arrow.utcnow()
     metrics.send('endpoint_updated', 'counter', 1, metric_tags={'source': endpoint.source.label})
     database.update(endpoint)
     return endpoint
diff --git a/lemur/sources/cli.py b/lemur/sources/cli.py
index bd72b343..77b9231f 100644
--- a/lemur/sources/cli.py
+++ b/lemur/sources/cli.py
@@ -56,7 +56,19 @@ def sync(source_strings):
         user = user_service.get_by_username('lemur')
 
         try:
-            source_service.sync(source, user)
+            data = source_service.sync(source, user)
+            print(
+                "[+] Certificates: New: {new} Updated: {updated}".format(
+                    new=data['certificates'][0],
+                    updated=data['certificates'][1]
+                )
+            )
+            print(
+                "[+] Endpoints: New: {new} Updated: {updated}".format(
+                    new=data['endpoints'][0],
+                    updated=data['endpoints'][1]
+                )
+            )
             print(
                 "[+] Finished syncing source: {label}. Run Time: {time}".format(
                     label=source.label,
diff --git a/lemur/sources/service.py b/lemur/sources/service.py
index cbe5654f..320c8985 100644
--- a/lemur/sources/service.py
+++ b/lemur/sources/service.py
@@ -10,7 +10,6 @@ import arrow
 from flask import current_app
 
 from lemur import database
-from lemur.extensions import metrics
 from lemur.sources.models import Source
 from lemur.certificates.models import Certificate
 from lemur.certificates import service as cert_service
@@ -45,24 +44,6 @@ def _disassociate_certs_from_source(certificates, source):
                 c.sources.delete(s)
 
 
-# TODO optimize via sql query
-def _disassociate_endpoints_from_source(endpoints, source):
-    current_endpoints = endpoint_service.get_by_source(source_label=source.label)
-
-    for ce in current_endpoints:
-        for fe in endpoints:
-            if ce.dnsname == fe['dnsname']:
-                break
-        else:
-            current_app.logger.info(
-                "Endpoint {dnsname} was not found during sync, removing from inventory.".format(
-                    dnsname=ce.dnsname
-                )
-            )
-            metrics.send('endpoint_removed', 'counter', 1)
-            database.delete(ce)
-
-
 def certificate_create(certificate, source):
     data, errors = CertificateUploadInputSchema().load(certificate)
 
@@ -121,7 +102,7 @@ def sync_endpoints(source):
             cert = cert_service.get_by_name(certificate_name)
 
         elif certificate:
-            cert = cert_service.get_by_body(certificate['body'])
+            cert = cert_service.find_duplicates(certificate)
             if not cert:
                 cert = cert_service.import_certificate(**certificate)
 
@@ -150,7 +131,7 @@ def sync_endpoints(source):
             endpoint_service.update(exists.id, **endpoint)
             updated += 1
 
-    _disassociate_endpoints_from_source(endpoints, source)
+    return new, updated
 
 
 def sync_certificates(source, user):
@@ -184,14 +165,18 @@ def sync_certificates(source, user):
     # we need to try and find the absent of certificates so we can properly disassociate them when they are deleted
     _disassociate_certs_from_source(certificates, source)
 
+    return new, updated
+
 
 def sync(source, user):
-    sync_certificates(source, user)
-    sync_endpoints(source)
+    new, updated = sync_certificates(source, user)
+    new, updated = sync_endpoints(source)
 
     source.last_run = arrow.utcnow()
     database.update(source)
 
+    return {'endpoints': (new, updated), 'certificates': (new, updated)}
+
 
 def clean(source):
     s = plugins.get(source.plugin_name)