From 848481fd1a738225664fa398275fcd1afd61c4ff Mon Sep 17 00:00:00 2001
From: demo <demo@antix1>
Date: Thu, 28 May 2026 17:28:59 -0400
Subject: feat: implement shortcode feature

---
 main.go | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

(limited to 'main.go')

diff --git a/main.go b/main.go
index 3582ed5..b48f6c9 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,7 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"fmt"
 	"log"
@@ -18,9 +19,11 @@ func main() {
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
 
 	maxConcurrency := flag.Int("c", 0, "Maximum number of concurrent queue pushes")
-	startRawURL := flag.String("url", "", "Entry-point URL")
+	urlArg := flag.String("url", "", "Entry-point URL")
 	maxURLs := flag.Int("max", 0, "Maximum number of URLs to collect (omitted or 0 means no limit)")
 	maxDepth := flag.Int("depth", 0, "Maximum URL depth (omitted or 0 means no limit)")
+	shortcode := flag.String("shortcode", "", "URL shortcode")
+	shortcodeFilename := flag.String("scfile", "urls.csv", "Shortcode CSV file")
 
 	flag.Parse()
 
@@ -34,8 +37,9 @@ func main() {
 		log.Fatalf("Invalid -c argument: %d", *maxConcurrency)
 	}
 
-	if *startRawURL == "" {
-		log.Fatal("Missing -url argument")
+	startRawURL, err := chooseFrom(*urlArg, *shortcode, *shortcodeFilename)
+	if err != nil {
+		log.Fatal(err)
 	}
 
 	if *maxURLs < 0 {
@@ -46,7 +50,7 @@ func main() {
 		log.Fatalf("Invalid -depth argument: %d", *maxDepth)
 	}
 
-	startURL, err := convertToURL(*startRawURL)
+	startURL, err := convertToURL(startRawURL)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -58,6 +62,24 @@ func main() {
 	})
 }
 
+// chooseFrom determines whether to use a -url or -shortcode
+// argument. If both are present or absent, an error is returned. If
+// exactly one is present, return that one.
+func chooseFrom(urlArg, shortcode, shortcodeFilename string) (string, error) {
+	urlAbsent := (urlArg == "")
+	shortcodeAbsent := (shortcode == "")
+
+	if urlAbsent == shortcodeAbsent {
+		return "", errors.New("-url and -shortcode flags either both missing or both present")
+	}
+
+	if urlAbsent {
+		return getURLFromShortcode(shortcodeFilename, shortcode)
+	}
+
+	return urlArg, nil
+}
+
 // convertToURL parses the given rawURL into a [url.URL]. If the
 // rawURL is missing a scheme, "https://" is prepended before parsing.
 //
-- 
cgit v1.2.3