From 03b8059a0f4bff3190b2b87037e3bfb0f65bdbe2 Mon Sep 17 00:00:00 2001
From: soulteary <soulteary@gmail.com>
Date: Wed, 15 Jun 2022 10:38:41 +0800
Subject: [PATCH] feat: support specify mirror by shorthand

---
 {linux => state}/custom.go | 20 +++++++++++++++++++-
 state/state.go             | 20 ++++++++++++++++++--
 2 files changed, 37 insertions(+), 3 deletions(-)
 rename {linux => state}/custom.go (73%)

diff --git a/linux/custom.go b/state/custom.go
similarity index 73%
rename from linux/custom.go
rename to state/custom.go
index 266ec4d..53576c9 100644
--- a/linux/custom.go
+++ b/state/custom.go
@@ -1,4 +1,4 @@
-package linux
+package state
 
 type buildin_custom_mirror struct {
 	url   string
@@ -22,3 +22,21 @@ var BUILDIN_CUSTOM_DEBIAN_MIRRORS = []buildin_custom_mirror{
 	{url: "https://repo.huaweicloud.com/debian/", alias: "cn:huawei"},
 	{url: "https://mirrors.cloud.tencent.com/debian/", alias: "cn:tencent"},
 }
+
+func getUbuntuMirrorByAliases(alias string) string {
+	for _, mirror := range BUILDIN_CUSTOM_UBUNTU_MIRRORS {
+		if mirror.alias == alias {
+			return mirror.url
+		}
+	}
+	return ""
+}
+
+func getDebianMirrorByAliases(alias string) string {
+	for _, mirror := range BUILDIN_CUSTOM_DEBIAN_MIRRORS {
+		if mirror.alias == alias {
+			return mirror.url
+		}
+	}
+	return ""
+}
diff --git a/state/state.go b/state/state.go
index 99503d3..18360d2 100644
--- a/state/state.go
+++ b/state/state.go
@@ -16,7 +16,15 @@ func GetProxyMode() int {
 	return PROXY_MODE
 }
 
-func SetUbuntuMirror(mirror string) {
+func SetUbuntuMirror(input string) {
+	mirror := ""
+	alias := getUbuntuMirrorByAliases(input)
+	if alias == "" {
+		mirror = input
+	} else {
+		mirror = alias
+	}
+
 	url, err := url.Parse(mirror)
 	if err != nil || mirror == "" {
 		UBUNTU_MIRROR = nil
@@ -33,7 +41,15 @@ func ResetUbuntuMirror() {
 	UBUNTU_MIRROR = nil
 }
 
-func SetDebianMirror(mirror string) {
+func SetDebianMirror(input string) {
+	mirror := ""
+	alias := getDebianMirrorByAliases(input)
+	if alias == "" {
+		mirror = input
+	} else {
+		mirror = alias
+	}
+
 	url, err := url.Parse(mirror)
 	if err != nil || mirror == "" {
 		DEBIAN_MIRROR = nil