From 5e84f98601d3c785df086df8a6e7821b82aebf4c Mon Sep 17 00:00:00 2001
From: Alderamin <126297468+Alderamin49ly@users.noreply.github.com>
Date: Fri, 19 Jul 2024 05:33:55 -0400
Subject: [PATCH] feat: add support for cmscript files (#1788)

Co-authored-by: Alderamin <alderamin49ly@gmail.com>
---
 pkg/tasks/heatmap.go | 38 +++++++++++++++++++++++---------------
 pkg/tasks/volume.go  |  9 +++++----
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/pkg/tasks/heatmap.go b/pkg/tasks/heatmap.go
index 50409d4d0..2ee89cbf3 100644
--- a/pkg/tasks/heatmap.go
+++ b/pkg/tasks/heatmap.go
@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
+	"strings"
 
 	"github.com/lucasb-eyer/go-colorful"
 	"github.com/sirupsen/logrus"
@@ -69,21 +70,24 @@ func GenerateHeatmaps(tlog *logrus.Entry) {
 				tlog.Infof("Generating heatmaps (%v/%v)", i+1, len(scriptfiles))
 			}
 			if file.Exists() {
-				log.Infof("Rendering %v", file.Filename)
-				destFile := filepath.Join(common.ScriptHeatmapDir, fmt.Sprintf("heatmap-%d.png", file.ID))
-				err := RenderHeatmap(
-					file.GetPath(),
-					destFile,
-					1000,
-					10,
-					250,
-				)
-				if err == nil {
-					file.HasHeatmap = true
-					file.RefreshHeatmapCache = true
-					file.Save()
-				} else {
-					log.Warn(err)
+				path := file.GetPath()
+				if strings.HasSuffix(path, ".funscript") {
+					log.Infof("Rendering %v", file.Filename)
+					destFile := filepath.Join(common.ScriptHeatmapDir, fmt.Sprintf("heatmap-%d.png", file.ID))
+					err := RenderHeatmap(
+						path,
+						destFile,
+						1000,
+						10,
+						250,
+					)
+					if err == nil {
+						file.HasHeatmap = true
+						file.RefreshHeatmapCache = true
+						file.Save()
+					} else {
+						log.Warn(err)
+					}
 				}
 			}
 		}
@@ -275,6 +279,10 @@ func (funscript Script) getDuration() float64 {
 }
 
 func getFunscriptDuration(path string) (float64, error) {
+	if !strings.HasSuffix(path, ".funscript") {
+		return 0.0, fmt.Errorf("Not a funscript: %s", path)
+	}
+
 	funscript, err := LoadFunscriptData(path)
 	if err != nil {
 		return 0.0, err
diff --git a/pkg/tasks/volume.go b/pkg/tasks/volume.go
index 4c9def118..085a7ded5 100644
--- a/pkg/tasks/volume.go
+++ b/pkg/tasks/volume.go
@@ -77,15 +77,16 @@ func RescanVolumes(id int) {
 			filename := escape(unescapedFilename)
 			filename2 := strings.Replace(filename, ".funscript", ".mp4", -1)
 			filename3 := strings.Replace(filename, ".hsp", ".mp4", -1)
-			filename3 = strings.Replace(filename3, ".srt", ".mp4", -1)
-			err := db.Where("filenames_arr LIKE ? OR filenames_arr LIKE ? OR filenames_arr LIKE ?", `%"`+filename+`"%`, `%"`+filename2+`"%`, `%"`+filename3+`"%`).Find(&scenes).Error
+			filename4 := strings.Replace(filename, ".srt", ".mp4", -1)
+			filename5 := strings.Replace(filename, ".cmscript", ".mp4", -1)
+			err := db.Where("filenames_arr LIKE ? OR filenames_arr LIKE ? OR filenames_arr LIKE ? OR filenames_arr LIKE ? OR filenames_arr LIKE ?", `%"`+filename+`"%`, `%"`+filename2+`"%`, `%"`+filename3+`"%`, `%"`+filename4+`"%`, `%"`+filename5+`"%`).Find(&scenes).Error
 			if err != nil {
 				log.Error(err, " when matching "+unescapedFilename)
 			}
 			if len(scenes) == 0 && config.Config.Advanced.UseAltSrcInFileMatching {
 				// check if the filename matches in external_reference record
 
-				db.Preload("XbvrLinks").Where("external_source like 'alternate scene %' and external_data LIKE ? OR external_data LIKE ? OR external_data LIKE ?", `%"`+filename+`%`, `%"`+filename2+`%`, `%"`+filename3+`%`).Find(&extrefs)
+				db.Preload("XbvrLinks").Where("external_source like 'alternate scene %' and external_data LIKE ? OR external_data LIKE ? OR external_data LIKE ? OR external_data LIKE ? OR external_data LIKE ?", `%"`+filename+`%`, `%"`+filename2+`%`, `%"`+filename3+`%`, `%"`+filename4+`%`, `%"`+filename5+`%`).Find(&extrefs)
 				if len(extrefs) == 1 {
 					if len(extrefs[0].XbvrLinks) == 1 {
 						// the scene id will be the Internal DB Id from the associated link
@@ -226,7 +227,7 @@ func scanLocalVolume(vol models.Volume, db *gorm.DB, tlog *logrus.Entry) {
 					}
 				}
 
-				if !strings.HasPrefix(filepath.Base(path), ".") && filepath.Ext(path) == ".funscript" {
+				if !strings.HasPrefix(filepath.Base(path), ".") && (filepath.Ext(path) == ".funscript" || filepath.Ext(path) == ".cmscript") {
 					scriptProcList = append(scriptProcList, path)
 				}
 				if !strings.HasPrefix(filepath.Base(path), ".") && filepath.Ext(path) == ".hsp" {