summaryrefslogtreecommitdiff
path: root/main.go
blob: d978717bc0965bc15f05e11a9a79b97f423ad821 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
/* SPDX-License-Identifier: MIT
 *
 * Copyright (C) 2020 jet tsang zeon-git. All Rights Reserved.
 */

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/antchfx/htmlquery"
	"github.com/valyala/fasthttp"
)

type ShortMap map[string]string

var listUrl string
var blockTxt string
var shortenMap map[string]string
var listen string
var reloadURI = "/reload"
var clientHeaderURI = "/header"
var UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.5060.53 Safari/537.36"

func init() {
	const (
		urlUsage     = "This program will get a list from the url you provide."
		blockUsage   = "html block to query."
		defaultBlock = ""

		defaultListen = "localhost:9530"
		listenUsage   = "TCP address to listen to"
	)
	flag.StringVar(&listUrl, "url", "https://pastebin.com/raw/9g4msZxC", urlUsage)
	flag.StringVar(&listUrl, "u", "https://pastebin.com/raw/9g4msZxC", urlUsage+" (shorthand)")

	flag.StringVar(&blockTxt, "block", defaultBlock, blockUsage)
	flag.StringVar(&blockTxt, "b", defaultBlock, blockUsage+" (shorthand)")

	flag.StringVar(&listen, "listen", defaultListen, listenUsage)
	flag.StringVar(&listen, "l", defaultListen, listenUsage+" (shorthand)")
	shortenMap = map[string]string{}
}
func main() {
	flag.Parse()
	go autoLoad()
	log.Println("Listening to:", listen)
	s := &fasthttp.Server{
		Handler:               requestHandler,
		NoDefaultServerHeader: true,
	}

	if err := s.ListenAndServe(listen); err != nil {
		log.Fatalf("Error in ListenAndServe: %s", err)
	}
}

func requestHandler(ctx *fasthttp.RequestCtx) {
	strr := string(ctx.RequestURI())
	ctx.Response.Header.Add("Server", "jet's relay server")
	if strr == reloadURI {
		loadData()
		fmt.Fprint(ctx, "ok")
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strr == clientHeaderURI {
		fmt.Fprint(ctx, string(ctx.Request.Header.Header()))
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strings.HasPrefix(strr, clientHeaderURI+"/") {
		headerkey := strings.TrimLeft(strr, clientHeaderURI+"/")
		log.Println("headerkey", headerkey)
		fmt.Fprint(ctx, string(ctx.Request.Header.Peek(headerkey)))
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strings.HasPrefix(strr, "/yt/") {
		var req fasthttp.Request
		var resp fasthttp.Response
		ythash := strings.TrimLeft(strr, "/yt/")
		req.Header.Add("User-Agent", UserAgent)
		req.Header.Add("Referer", "https://www.youtube.com")
		req.SetRequestURI("https://www.youtube.com/watch?v=" + ythash)
		fasthttp.Do(&req, &resp)
		ytbody := string(resp.Body())
		m3u8 := GetBetweenStr(ytbody, "hlsManifestUrl\":\"", "\"},\"heart")
		ctx.Redirect(m3u8, 302)
		defer func() {
			fasthttp.ReleaseRequest(&req)
			req.ConnectionClose()
			resp.ConnectionClose()
		}()
		return
	}
	if relayUrl, exist := shortenMap[string(ctx.Path())]; exist {
		var req fasthttp.Request
		ctx.Request.Header.CopyTo(&req.Header)
		req.SetRequestURI(relayUrl)
		req.Header.SetBytesV("Host", req.Host())
		ctx.URI().QueryArgs().VisitAll(func(key []byte, value []byte) {
			req.URI().QueryArgs().AddBytesKV(key, value)
		})
		defer fasthttp.ReleaseRequest(&req)

		fasthttp.Do(&req, &ctx.Response)
	} else {
		ctx.Error("err", 404)
	}
}

func autoLoad() {
	for {
		loadData()
		time.Sleep(time.Hour)
	}
}

func loadData() {
	log.Println("Loading data.", listUrl)
	var code string

	if blockTxt != "" {
		doc, err := htmlquery.LoadURL(listUrl)
		if err != nil {
			log.Println("url err:", listUrl)
			return
		}
		found, err := htmlquery.QueryAll(doc, blockTxt)
		if err != nil {
			log.Println("block not found", err)
			return
		}
		if len(found) == 0 {
			log.Println("No elements found")
			return
		}
		code = htmlquery.InnerText(found[0])
	} else {
		var dst []byte
		statusCode, byteCode, err := fasthttp.Get(dst, string(listUrl))
		if err != nil || statusCode != fasthttp.StatusOK {
			log.Printf("HTTP request failed with status code %d\n", statusCode)
			return
		}
		code = string(byteCode)
	}

	scanner := bufio.NewScanner(strings.NewReader(code))
	for scanner.Scan() {
		maplist := strings.Split(scanner.Text(), "#")
		if len(maplist) < 2 {
			log.Println("Invalid data format:", scanner.Text())
			continue // Skip this iteration as the data format is not valid
		}
		shortenMap[maplist[0]] = maplist[1]
	}
	log.Println("URLs mapped.")
}

func GetBetweenStr(str, start, end string) string {
	n := strings.Index(str, start) + len(start)
	if n == -1 {
		n = 0
	}
	str = string([]byte(str)[n:])
	m := strings.Index(str, end)
	if m == -1 {
		m = len(str)
	}
	str = string([]byte(str)[:m])
	return str
}