From a9bef5d86896dd646cf3496984a8770e4764e387 Mon Sep 17 00:00:00 2001
From: Todd Bednarczyk <todd.bednarczyk@gmail.com>
Date: Thu, 2 Jan 2020 05:13:29 -0500
Subject: [PATCH] Add TikTok support (#591)

* added tiktok support

* Added tiktok to readme
---
 README.md                        |  1 +
 extractors/tiktok/tiktok.go      | 61 ++++++++++++++++++++++++++++++++
 extractors/tiktok/tiktok_test.go | 37 +++++++++++++++++++
 main.go                          |  3 ++
 4 files changed, 102 insertions(+)
 create mode 100644 extractors/tiktok/tiktok.go
 create mode 100644 extractors/tiktok/tiktok_test.go

diff --git a/README.md b/README.md
index 5169a32e1..6ae40c954 100644
--- a/README.md
+++ b/README.md
@@ -615,6 +615,7 @@ Twitter | <https://twitter.com> | ✓ | | | |
 Pornhub | <https://pornhub.com> | ✓ | | | |
 XVIDEOS | <https://xvideos.com> | ✓ | | | |
 聯合新聞網 | <https://udn.com> | ✓ | | | |
+TikTok | <https://www.tiktok.com> | ✓ | | | |
 
 
 ## Known issues
diff --git a/extractors/tiktok/tiktok.go b/extractors/tiktok/tiktok.go
new file mode 100644
index 000000000..413ae076d
--- /dev/null
+++ b/extractors/tiktok/tiktok.go
@@ -0,0 +1,61 @@
+package tiktok
+
+import (
+	"github.com/iawia002/annie/downloader"
+	"github.com/iawia002/annie/extractors"
+	"github.com/iawia002/annie/request"
+	"github.com/iawia002/annie/utils"
+)
+
+// Extract is the main function for extracting data
+func Extract(uri string) ([]downloader.Data, error) {
+	html, err := request.Get(uri, uri, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// There are a few json objects loaded into the html that are useful. We're able to parse the video url from the
+	// videoObject json.
+
+	videoScriptTag := utils.MatchOneOf(html, `<script type="application\/ld\+json" id="videoObject">(.*?)<\/script>`)
+	if videoScriptTag == nil || len(videoScriptTag) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	videoJSON := videoScriptTag[1]
+	videoURL := utils.GetStringFromJson(videoJSON, "contentUrl")
+
+	// We can receive the title directly from this __NEXT_DATA__ object.
+
+	nextScriptTag := utils.MatchOneOf(html, `<script id="__NEXT_DATA__" type="application\/json" crossorigin="anonymous">(.*?)<\/script>`)
+	if nextScriptTag == nil || len(nextScriptTag) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	nextJSON := nextScriptTag[1]
+	title := utils.GetStringFromJson(nextJSON, "props.pageProps.videoData.itemInfos.text")
+
+	streams := map[string]downloader.Stream{}
+
+	size, err := request.Size(videoURL, uri)
+	if err != nil {
+		return nil, err
+	}
+	urlData := downloader.URL{
+		URL:  videoURL,
+		Size: size,
+		Ext:  "mp4",
+	}
+	streams["default"] = downloader.Stream{
+		URLs: []downloader.URL{urlData},
+		Size: size,
+	}
+
+	return []downloader.Data{
+		{
+			Site:    "TikTok tiktok.com",
+			Title:   title,
+			Type:    "video",
+			Streams: streams,
+			URL:     uri,
+		},
+	}, nil
+}
diff --git a/extractors/tiktok/tiktok_test.go b/extractors/tiktok/tiktok_test.go
new file mode 100644
index 000000000..a5dedb2cf
--- /dev/null
+++ b/extractors/tiktok/tiktok_test.go
@@ -0,0 +1,37 @@
+package tiktok
+
+import (
+	"github.com/iawia002/annie/config"
+	"github.com/iawia002/annie/test"
+	"testing"
+)
+
+func TestDownload(t *testing.T) {
+	config.InfoOnly = true
+	tests := []struct {
+		name string
+		args test.Args
+	}{
+		{
+			name: "normal test",
+			args: test.Args{
+				URL:   "https://www.tiktok.com/@therock/video/6768158408110624005",
+				Title: "#bestfriend check.",
+			},
+		},
+		{
+			name: "short url test",
+			args: test.Args{
+				URL:   "https://vm.tiktok.com/C998PY/",
+				Title: "Who saw that coming? 🍁 #leaves #fall",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			data, err := Extract(tt.args.URL)
+			test.CheckError(t, err)
+			test.Check(t, tt.args, data[0])
+		})
+	}
+}
diff --git a/main.go b/main.go
index 0e3f18a80..814ac25b2 100644
--- a/main.go
+++ b/main.go
@@ -27,6 +27,7 @@ import (
 	"github.com/iawia002/annie/extractors/pornhub"
 	"github.com/iawia002/annie/extractors/qq"
 	"github.com/iawia002/annie/extractors/tangdou"
+	"github.com/iawia002/annie/extractors/tiktok"
 	"github.com/iawia002/annie/extractors/tumblr"
 	"github.com/iawia002/annie/extractors/twitter"
 	"github.com/iawia002/annie/extractors/udn"
@@ -163,6 +164,8 @@ func download(videoURL string) bool {
 		data, err = xvideos.Extract(videoURL)
 	case "udn":
 		data, err = udn.Extract(videoURL)
+	case "tiktok":
+		data, err = tiktok.Extract(videoURL)
 	default:
 		data, err = universal.Extract(videoURL)
 	}