ugoira-tool/main.go
2024-12-31 14:12:28 -03:00

206 lines
5.6 KiB
Go

package main
import (
"flag"
"fmt"
"io"
"net/http"
"os"
"strconv"
)
type CliFlags struct {
outputFileName string
manifestFileName string
userAgent string
ffmpegArgs string
verbose bool
}
type PostId uint64
const (
ArtworkUrlTemplate = "https://www.pixiv.net/en/artworks/%d"
ManifestUrlTemplate = "https://www.pixiv.net/ajax/illust/%d/ugoira_meta"
)
func doPixivRequest(client *http.Client, url, referer, userAgent string) io.ReadCloser {
req, err := http.NewRequest("GET", url, nil)
if err != nil {
die("couldn't create request: %s", err)
}
req.Header.Set("Authority", "www.pixiv.net")
req.Header.Set("Sec-Fetch-Dest", "empty")
req.Header.Set("Sec-Fetch-Mode", "cors")
req.Header.Set("Sec-Fetch-Site", "same-origin")
req.Header.Set("Referer", referer)
if len(userAgent) > 0 {
req.Header.Set("User-Agent", userAgent)
}
response, err := client.Do(req)
if err != nil {
die("GET request failed: %s", err)
}
if response.StatusCode != http.StatusOK {
die("Got non-200 status code: %d", response.StatusCode)
}
return response.Body
}
func getUgoira(postId PostId, flags CliFlags) {
client := &http.Client{}
artworkUrl := fmt.Sprintf(ArtworkUrlTemplate, postId)
manifestUrl := fmt.Sprintf(ManifestUrlTemplate, postId)
say(flags.verbose, "Fetching and decoding manifest")
manifestData := doPixivRequest(client, manifestUrl, artworkUrl, flags.userAgent)
defer manifestData.Close()
// Read the entire JSON in memory, since we need to dump it to a file later
manifestJsonData, err := io.ReadAll(manifestData)
if err != nil {
die("error while fetching manifest: %s", err)
}
// Now decode the JSON...
manifest, err := decodeUgoiraManifest(manifestJsonData)
if err != nil {
die("%s", err)
}
if manifest.Error {
die("Manifest returned error: %s", manifest.Message)
}
if len(manifest.Body.OriginalSrc) == 0 {
die("Ugoira url is empty")
}
say(flags.verbose, "Ugoira file URL: %s", manifest.Body.OriginalSrc)
// Download ugoira data
var ugoiraFileName string
if len(flags.outputFileName) > 0 {
ugoiraFileName = flags.outputFileName
} else {
ugoiraFileName = fmt.Sprintf("./%d_ugoira.zip", postId)
}
say(flags.verbose, "Downloading ugoira (saving to: %s)", ugoiraFileName)
ugoiraData := doPixivRequest(client, manifest.Body.OriginalSrc, artworkUrl, flags.userAgent)
ugoiraFile, err := os.Create(ugoiraFileName)
if err != nil {
die("Could not create ugoira file: %s", err)
}
defer ugoiraFile.Close()
if _, err := io.Copy(ugoiraFile, ugoiraData); err != nil {
die("Ugoira download failed: %s", err)
}
// Dump manifest JSON
var ugoiraManifestFileName string
if len(flags.manifestFileName) > 0 {
ugoiraManifestFileName = flags.manifestFileName
} else {
ugoiraManifestFileName = fmt.Sprintf("./%d_ugoira.json", postId)
}
say(flags.verbose, "Dumping manifest (to: %s)", ugoiraManifestFileName)
ugoiraManifestFile, err := os.Create(ugoiraManifestFileName)
if err != nil {
die("Could not create ugoira manifest file: %s", err)
}
defer ugoiraManifestFile.Close()
if _, err := ugoiraManifestFile.Write(manifestJsonData); err != nil {
die("Ugoira manifest dump failed: %s", err)
}
say(flags.verbose, "Successfully downloaded post %d", postId)
}
func die(format string, args ...any) {
fmt.Fprintf(os.Stderr, os.Args[0]+": "+format+"\n", args...)
os.Exit(1)
}
func say(enabled bool, format string, args ...any) {
if !enabled {
return
}
fmt.Printf(format+"\n", args...)
}
func downloadUgoiraCmd(args []string, flags CliFlags) {
if len(args) == 0 {
die("expected post ID")
}
postIdStr := args[0]
postId, err := strconv.ParseUint(postIdStr, 10, 0)
if err != nil {
die("Invalid post ID (has to be numerical): '%s'", postIdStr)
}
say(flags.verbose, "Working on post %d", postId)
getUgoira(PostId(postId), flags)
}
func ugoira2videoCmd(args []string, flags CliFlags) {
if len(args) < 2 {
die("expected ugoira zip file and manifest JSON")
}
ugoiraFileName, manifestFileName := args[0], args[1]
mfp, err := os.Open(manifestFileName)
if err != nil {
die("Could not open manifest file: %s", err)
}
defer mfp.Close()
manifestData, err := io.ReadAll(mfp)
if err != nil {
die("Could not read manifest: %s", err)
}
manifest, err := decodeUgoiraManifest(manifestData)
if err != nil {
die("%s", err)
}
err = ugoira2video(manifest, ugoiraFileName, flags.outputFileName, flags.ffmpegArgs)
if err != nil {
die("Conversion failed: %s", err)
}
}
func main() {
var flags CliFlags
downloadCmd := flag.NewFlagSet("download", flag.ExitOnError)
downloadCmd.StringVar(&flags.outputFileName, "o", "", "Ugoira output file name")
downloadCmd.StringVar(&flags.manifestFileName, "M", "", "Ugoira manifest output file name")
downloadCmd.StringVar(&flags.userAgent, "A", "", "User agent to use when doing GET requests")
downloadCmd.BoolVar(&flags.verbose, "v", false, "Enable verbosity")
u2vCmd := flag.NewFlagSet("u2v", flag.ExitOnError)
u2vCmd.StringVar(&flags.outputFileName, "o", "", "Converted Ugoira output file")
u2vCmd.StringVar(
&flags.ffmpegArgs,
"ffmpeg-args",
"",
"Arguments to be supplied to FFmpeg (a single `%s` acts as a placeholder for the output file name)",
)
if len(os.Args) < 2 {
die("expected command")
}
switch os.Args[1] {
case "download", "down", "d":
downloadCmd.Parse(os.Args[2:])
downloadUgoiraCmd(downloadCmd.Args(), flags)
case "u2v":
u2vCmd.Parse(os.Args[2:])
if len(flags.outputFileName) == 0 {
die("Expected filename output for video")
}
ugoira2videoCmd(u2vCmd.Args(), flags)
case "-h":
fmt.Printf("Commands: download, u2v\nUse '%s <command> -h' to see the command-specific flags.\n", os.Args[0])
default:
die("invalid command '%s'", os.Args[1])
}
}