pixiv-scrapper/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"path"
	"syscall"

	"git.magicalmirro.red/voice/pixiv-scrapper/pixiv"
)

type cmdlineT struct {
	followings    bool
	favs          bool
	illusts       bool
	userID        string
	setCookies    string
	setProxy      string
	workDirectory string
	userAgent     string
	printID       bool
	printURL      bool
	test          bool
	logErrorsFile string
	threads       int
	setXAttrTags  bool
	setTagsFile   string
	testString    string
}

var (
	cmdline cmdlineT
	cookies string
	pix     pixiv.Pixiv
)

func main() {
	flags()
	pix = pixiv.New(cmdline.setCookies, cmdline.logErrorsFile, cmdline.threads, cmdline.setXAttrTags)
	defer pix.Close()
	if len(cmdline.setProxy) > 0 {
		pix.SetProxy(cmdline.setProxy)
	}
	if len(cmdline.workDirectory) > 0 {
		pix.WorkDirectory = cmdline.workDirectory
	}

	if cmdline.illusts {
		err := fetchIllusts(cmdline.userID)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	if cmdline.favs {
		err := fetchBookmarks(cmdline.userID)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	if cmdline.followings {
		err := fetchFollows(cmdline.userID)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	if len(cmdline.setTagsFile) > 0 {
		err := setTagsFile(cmdline.setTagsFile)
		if err != nil {
			fmt.Printf("%s\n", err.Error())
		}
	}
}

func fetchFollows(userID string) (err error) {
	users, err := pix.GetFollows(userID)
	if err != nil {
		return
	}
	for _, user := range users {
		fmt.Printf("%s %s %s %s\n", user.UserID, "followed_user", "defered", user.UserName)
	}
	return
}

func fetchBookmarks(userID string) (err error) {
	illusts, err := pix.GetBookmarks(userID)
	if err != nil {
		return
	}
	for _, illust := range illusts {
		if cmdline.printID {
			fmt.Printf("%s %s %s %s %s\n", illust.ID, "illust", "defered", illust.Title, illust.Alt)
			continue
		}
		if !illust.Complited() {
			err := pix.ComplateIllust(&illust)
			if err != nil {
				return err
			}
		}
		if cmdline.printURL {
			for _, page := range illust.Pages {
				fmt.Println(page.URLs.Original)
			}
			continue
		}
		if !(cmdline.printURL && cmdline.printID) {
			if !(cmdline.printURL && cmdline.printID) {
				for _, illust := range illusts {
					pix.DownloadChannel <- illust
				}
			}
		}

	}
	return
}

func setTagsFile(filename string) (err error) {
	if ok, id := pix.TypicalFilenamesMatcher(path.Base(filename)); ok {
		illust, err := pix.GetIllust(id)
		if err != nil {
			return err
		}
		err = syscall.Setxattr(filename, "user.xdg.tags", []byte(illust.TagsString()), 0)
		if err != nil {
			return err
		}
	}
	return
}

func fetchIllusts(userID string) (err error) {
	illusts, err := pix.GetUserIllustsID(userID)
	if err != nil {
		return
	}
	if cmdline.printID {
		for _, illust := range illusts {
			fmt.Println(illust.ID)
		}
		return
	}
	if cmdline.printURL {
		for _, illust := range illusts {
			id := illust.ID
			illust, err := pix.GetIllust(id)
			if err != nil {
				return err
			}
			for _, page := range illust.Pages {
				fmt.Println(page.URLs.Original)
			}
		}
		return
	}
	if !(cmdline.printURL && cmdline.printID) {
		for _, illust := range illusts {
			pix.DownloadChannel <- illust
		}
	}

	if len(cmdline.setTagsFile) > 0 {
		err = setTagsFile(cmdline.setTagsFile)
		if err != nil {
			fmt.Printf("Request error or ID not recognized: %s\n", err.Error())
		}
	}
	return
}

func flags() {
	flag.BoolVar(&cmdline.illusts, "illusts", false, "download illusts")
	flag.BoolVar(&cmdline.followings, "followings", false, "fetch user followings")
	flag.BoolVar(&cmdline.favs, "favs", false, "fetch user favorites")
	flag.StringVar(&cmdline.userID, "user-id", "", "set userID to extract followings, favs or illustrations from")
	flag.StringVar(&cmdline.setCookies, "set-cookies", "", "set cookies for request. PHPSESSID is REQUIRED if you want to see R-18 works. SHOULD be something like PHPSESSID=12345678_MEPujYEEwOfXKYdB9TXeMIxaq7pQPYXA")
	flag.StringVar(&cmdline.setProxy, "set-proxy", "", "set proxy for request")
	flag.StringVar(&cmdline.workDirectory, "workdir", "", "set directory to save images in")
	flag.StringVar(&cmdline.userAgent, "user-agent", "Mozilla/5.0 (Windows NT 10.0; rv:78.0) Gecko/20100101 Firefox/78.0", "set user agent")
	flag.BoolVar(&cmdline.printID, "print-id", false, "only print IDs, not download")
	flag.BoolVar(&cmdline.printURL, "print-url", false, "only print URLs, not download")
	flag.BoolVar(&cmdline.test, "test", false, "test")
	flag.StringVar(&cmdline.logErrorsFile, "log-errors", "", "file to strore failed items")
	flag.IntVar(&cmdline.threads, "threads", 1, "threads number")
	flag.StringVar(&cmdline.setTagsFile, "set-xattr-file", "", "fetch illustration's tags and write them to xattrs of the specified file.")
	flag.BoolVar(&cmdline.setXAttrTags, "set-xattr", false, "also write tags to xattrs")
	flag.StringVar(&cmdline.testString, "test-string", "", "used for testing purposes")
	flag.Parse()
}