From 2773927090590c1fe7d3e485998cf204879452c3 Mon Sep 17 00:00:00 2001
From: Ice3man <nizamulrana@gmail.com>
Date: Fri, 25 Aug 2023 17:59:56 +0530
Subject: [PATCH] feat: added template-url support in template flag feature

---
 v2/pkg/catalog/loader/loader.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/v2/pkg/catalog/loader/loader.go b/v2/pkg/catalog/loader/loader.go
index 8535b82334..fb9c23eb78 100644
--- a/v2/pkg/catalog/loader/loader.go
+++ b/v2/pkg/catalog/loader/loader.go
@@ -1,8 +1,11 @@
 package loader
 
 import (
+	"fmt"
+	"net/url"
 	"os"
 	"sort"
+	"strings"
 
 	"github.com/pkg/errors"
 	"github.com/projectdiscovery/gologger"
@@ -120,6 +123,20 @@ func New(config *Config) (*Store, error) {
 		finalWorkflows: config.Workflows,
 	}
 
+	// Do a check to see if we have URLs in templates flag, if so
+	// we need to processs them separately and remove them from the initial list
+	var templatesFinal []string
+	var urlTemplates []string
+	for _, template := range config.Templates {
+		if _, err := url.ParseRequestURI(template); err == nil {
+			urlTemplates = append(urlTemplates, handleTemplatesEditorURLs(template))
+		} else {
+			templatesFinal = append(templatesFinal, template)
+		}
+	}
+	store.finalTemplates = templatesFinal
+	config.TemplateURLs = append(config.TemplateURLs, urlTemplates...)
+
 	urlBasedTemplatesProvided := len(config.TemplateURLs) > 0 || len(config.WorkflowURLs) > 0
 	if urlBasedTemplatesProvided {
 		remoteTemplates, remoteWorkflows, err := getRemoteTemplatesAndWorkflows(config.TemplateURLs, config.WorkflowURLs, config.RemoteTemplateDomainList)
@@ -145,6 +162,22 @@ func New(config *Config) (*Store, error) {
 	return store, nil
 }
 
+func handleTemplatesEditorURLs(input string) string {
+	parsed, err := url.Parse(input)
+	if err != nil {
+		return input
+	}
+	if !strings.HasSuffix(parsed.Hostname(), "ondigitalocean.app") {
+		return input
+	}
+	if strings.HasSuffix(parsed.Path, ".yaml") {
+		return input
+	}
+	parsed.Path = fmt.Sprintf("%s.yaml", parsed.Path)
+	finalURL := parsed.String()
+	return finalURL
+}
+
 // Templates returns all the templates in the store
 func (store *Store) Templates() []*templates.Template {
 	return store.templates