From 012ac6402f8a2b6f5dc49272596e0b5408609330 Mon Sep 17 00:00:00 2001
From: afabiani <alessio.fabiani@geo-solutions.it>
Date: Fri, 15 Nov 2019 15:28:39 +0100
Subject: [PATCH] [Hardening] fixup Legend links

---
 .../commands/set_all_layers_metadata.py          | 16 ++++++++++++++++
 geonode/base/models.py                           |  4 ++--
 .../layers/templates/layers/layer_detail.html    |  2 +-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/geonode/base/management/commands/set_all_layers_metadata.py b/geonode/base/management/commands/set_all_layers_metadata.py
index 2eb9334a6df..42d4b84277d 100644
--- a/geonode/base/management/commands/set_all_layers_metadata.py
+++ b/geonode/base/management/commands/set_all_layers_metadata.py
@@ -18,6 +18,7 @@
 #
 #########################################################################
 
+from geonode.geoserver.helpers import ogc_server_settings
 from django.core.management.base import BaseCommand
 from geonode.base.models import Link
 from geonode.layers.models import Layer
@@ -104,6 +105,21 @@ def handle(self, *args, **options):
                         while _links.count() > 1:
                             _links.last().delete()
                             print '.',
+                    # fixup Legend links
+                    legend_url_template = ogc_server_settings.PUBLIC_LOCATION + \
+                        'ows?service=WMS&request=GetLegendGraphic&format=image/png&WIDTH=20&HEIGHT=20&LAYER=' + \
+                        '{alternate}&STYLE={style_name}' + \
+                        '&legend_options=fontAntiAliasing:true;fontSize:12;forceLabels:on'
+                    if not layer.get_legend_url(style_name=layer.default_style.name):
+                        Link.objects.update_or_create(
+                            resource=layer.resourcebase_ptr,
+                            name='Legend',
+                            extension='png',
+                            url=legend_url_template.format(
+                                alternate=layer.alternate,
+                                style_name=layer.default_style.name),
+                            mime='image/png',
+                            link_type='image')
             except BaseException as e:
                 import traceback
                 traceback.print_exc()
diff --git a/geonode/base/models.py b/geonode/base/models.py
index 67a22837087..2a263f8ef27 100644
--- a/geonode/base/models.py
+++ b/geonode/base/models.py
@@ -1107,12 +1107,12 @@ def get_legend_url(self, style_name=None):
 
         if legend.count() > 0:
             if not style_name:
-                return legend[0].url
+                return legend.first().url
             else:
                 for _legend in legend:
                     if style_name in _legend.url:
                         return _legend.url
-        return legend.url
+        return None
 
     def get_ows_url(self):
         """Return URL for OGC WMS server None if it does not exist.
diff --git a/geonode/layers/templates/layers/layer_detail.html b/geonode/layers/templates/layers/layer_detail.html
index d0686d7e890..2f009cbfa2a 100644
--- a/geonode/layers/templates/layers/layer_detail.html
+++ b/geonode/layers/templates/layers/layer_detail.html
@@ -580,7 +580,7 @@ <h5 class="modal-title" id="myModalLabel">{% trans "Attribute Information" %}</h
       <h4 class="list-group-item-heading">{%trans "Legend" %}</h4>
       {% if resource.default_style %}
       {% for style in resource.styles.all %}
-          {% if resource.default_style == style %}
+          {% if resource.default_style == style or resource.default_style.name == style.name or resource.default_style.sld_title == style.sld_title %}
               {% for legend in resource.get_legend %}
                   {% with "STYLE="|add:style.name as style_name %}
                   {% if legend.link_type == 'image' and style_name in legend.url %}