379 lines
12 KiB
Go
379 lines
12 KiB
Go
package e621
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/gzip"
|
|
"encoding/csv"
|
|
"fmt"
|
|
"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/builder"
|
|
"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/model"
|
|
"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/utils"
|
|
"github.com/gocarina/gocsv"
|
|
_ "github.com/joho/godotenv/autoload"
|
|
"golang.org/x/time/rate"
|
|
"log"
|
|
"math"
|
|
"net/http"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// Client is the main client for interacting with the e621 API.
|
|
type Client struct {
|
|
RequestContext model.RequestContext
|
|
}
|
|
|
|
// NewClient creates a new e621 client with the provided username and API key.
|
|
func NewClient(username string, apiKey string) Client {
|
|
// Create a new e621 client with the given username and API key.
|
|
return Client{
|
|
RequestContext: model.RequestContext{
|
|
Client: http.Client{},
|
|
RateLimiter: rate.NewLimiter(1, 2),
|
|
Host: "https://e621.net",
|
|
UserAgent: fmt.Sprintf("Go-e621-SDK used by %s | (made by the Anthrove Team)", username),
|
|
Username: username,
|
|
APIKey: apiKey,
|
|
},
|
|
}
|
|
}
|
|
|
|
// GetUserBuilder returns a UserBuilder instance for creating and executing requests to retrieve user information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.UserBuilder: An instance of the UserBuilder.
|
|
func (c *Client) GetUserBuilder() builder.UserBuilder {
|
|
return builder.NewGetUserBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetUsersBuilder returns a UsersBuilder instance for creating and executing requests to retrieve multiple users' information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.UsersBuilder: An instance of the UsersBuilder.
|
|
func (c *Client) GetUsersBuilder() builder.UsersBuilder {
|
|
return builder.NewGetUsersBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetFavoritesBuilder returns a FavoritesBuilder instance for creating and executing requests to retrieve a user's favorite posts from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.FavoritesBuilder: An instance of the FavoritesBuilder.
|
|
func (c *Client) GetFavoritesBuilder() builder.FavoritesBuilder {
|
|
return builder.NewGetFavoritesBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetPostBuilder returns a PostBuilder instance for creating and executing requests to retrieve post information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.PostBuilder: An instance of the PostBuilder.
|
|
func (c *Client) GetPostBuilder() builder.PostBuilder {
|
|
return builder.NewGetPostBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetPostsBuilder returns a PostsBuilder instance for creating and executing requests to retrieve multiple posts' information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.PostsBuilder: An instance of the PostsBuilder.
|
|
func (c *Client) GetPostsBuilder() builder.PostsBuilder {
|
|
return builder.NewGetPostsBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetNoteBuilder returns a NoteBuilder instance for creating and executing requests to retrieve note information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.NoteBuilder: An instance of the NoteBuilder.
|
|
func (c *Client) GetNoteBuilder() builder.NoteBuilder {
|
|
return builder.NewGetNoteBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetNotesBuilder returns a NotesBuilder instance for creating and executing requests to retrieve multiple notes' information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.NotesBuilder: An instance of the NotesBuilder.
|
|
func (c *Client) GetNotesBuilder() builder.NotesBuilder {
|
|
return builder.NewGetNotesBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetPoolBuilder returns a PoolBuilder instance for creating and executing requests to retrieve pool information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.PoolBuilder: An instance of the PoolBuilder.
|
|
func (c *Client) GetPoolBuilder() builder.PoolBuilder {
|
|
return builder.NewGetPoolBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetPoolsBuilder returns a PoolsBuilder instance for creating and executing requests to retrieve multiple pools' information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.PoolsBuilder: An instance of the PoolsBuilder.
|
|
func (c *Client) GetPoolsBuilder() builder.PoolsBuilder {
|
|
return builder.NewGetPoolsBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetTagBuilder returns a TagBuilder instance for creating and executing requests to retrieve tag information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.TagBuilder: An instance of the TagBuilder.
|
|
func (c *Client) GetTagBuilder() builder.TagBuilder {
|
|
return builder.NewGetTagBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetTagsBuilder returns a TagsBuilder instance for creating and executing requests to retrieve multiple tags' information from the e621 API.
|
|
//
|
|
// Returns:
|
|
// - builder.TagsBuilder: An instance of the TagsBuilder.
|
|
func (c *Client) GetTagsBuilder() builder.TagsBuilder {
|
|
return builder.NewGetTagsBuilder(c.RequestContext)
|
|
}
|
|
|
|
// SetHost sets the API host for the client.
|
|
func (c *Client) SetHost(host string) *Client {
|
|
// Set the API host for the client.
|
|
c.RequestContext.Host = host
|
|
return c
|
|
}
|
|
|
|
// SetAgentName sets the user agent name for the client.
|
|
func (c *Client) SetAgentName(userAgent string) *Client {
|
|
// Set the user agent name for the client.
|
|
c.RequestContext.UserAgent = userAgent
|
|
return c
|
|
}
|
|
|
|
// GetUserByName returns a user builder for a given username.
|
|
func (c *Client) GetUserByName(username string) builder.UserBuilder {
|
|
// Returns a user builder for the specified username.
|
|
return builder.NewGetUserBuilder(c.RequestContext).SetUsername(username)
|
|
}
|
|
|
|
// GetUserByID returns a user builder for a given user ID.
|
|
func (c *Client) GetUserByID(userID model.UserID) builder.UserBuilder {
|
|
// Returns a user builder for the specified user ID.
|
|
return builder.NewGetUserBuilder(c.RequestContext).SetUsername(strconv.FormatInt(int64(userID), 10))
|
|
}
|
|
|
|
// GetFavoritesForUser returns a favorites builder for a given username.
|
|
func (c *Client) GetFavoritesForUser(username string) (builder.FavoritesBuilder, error) {
|
|
// Returns a favorites builder for the specified username.
|
|
user, err := builder.NewGetUserBuilder(c.RequestContext).SetUsername(username).Execute()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
favoritesBuilder := builder.NewGetFavoritesBuilder(c.RequestContext).SetUserID(user.ID)
|
|
return favoritesBuilder, nil
|
|
}
|
|
|
|
// GetNFavoritesForUser retrieves a specified number of favorites for a user.
|
|
func (c *Client) GetNFavoritesForUser(n int, favoriteBuilder builder.FavoritesBuilder) ([]model.Post, error) {
|
|
// Retrieves a specified number of favorite posts for a user.
|
|
if n < utils.E621_MAX_POST_COUNT {
|
|
favoriteBuilder.SetLimit(n)
|
|
}
|
|
|
|
var favorites []model.Post
|
|
var page = 1
|
|
for len(favorites) < n {
|
|
favoriteBuilder.Page(page)
|
|
favoriteBuilder.SetLimit(n - len(favorites))
|
|
newFavorites, err := favoriteBuilder.Execute()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(newFavorites) == 0 {
|
|
break
|
|
}
|
|
favorites = append(favorites, newFavorites...)
|
|
page = page + 1
|
|
}
|
|
|
|
return favorites, nil
|
|
}
|
|
|
|
// GetAllFavoritesForUser retrieves all favorites for a user.
|
|
func (c *Client) GetAllFavoritesForUser(favoriteBuilder builder.FavoritesBuilder) ([]model.Post, error) {
|
|
// Retrieves all favorite posts for a user.
|
|
return c.GetNFavoritesForUser(math.MaxInt, favoriteBuilder)
|
|
}
|
|
|
|
// GetFavoritesForUserWithTags returns a posts builder for a user's favorites with specific tags.
|
|
func (c *Client) GetFavoritesForUserWithTags(username string, tags string) builder.PostsBuilder {
|
|
// Returns a posts builder for a user's favorites with specific tags.
|
|
favoritesBuilder := builder.NewGetPostsBuilder(c.RequestContext).Tags(fmt.Sprintf("fav:%s %s", username, tags))
|
|
return favoritesBuilder
|
|
}
|
|
|
|
// GetPostByID returns a post builder for a specific post ID.
|
|
func (c *Client) GetPostByID(id model.PostID) builder.PostBuilder {
|
|
// Returns a post builder for a specific post ID.
|
|
return builder.NewGetPostBuilder(c.RequestContext).SetPostID(id)
|
|
}
|
|
|
|
// GetPosts returns a posts builder for general post queries.
|
|
func (c *Client) GetPosts() builder.PostsBuilder {
|
|
// Returns a posts builder for general post queries.
|
|
return builder.NewGetPostsBuilder(c.RequestContext)
|
|
}
|
|
|
|
// GetNPosts retrieves a specified number of posts.
|
|
func (c *Client) GetNPosts(n int, postBuilder builder.PostsBuilder) ([]model.Post, error) {
|
|
// Retrieves a specified number of posts using the provided post builder.
|
|
if n < utils.E621_MAX_POST_COUNT {
|
|
postBuilder.SetLimit(n)
|
|
}
|
|
|
|
posts, err := postBuilder.Execute()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
for len(posts) < n {
|
|
lastPostID := posts[len(posts)-1].ID
|
|
postBuilder.PageBefore(lastPostID)
|
|
postBuilder.SetLimit(n - len(posts))
|
|
newPosts, err := postBuilder.Execute()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(newPosts) == 0 {
|
|
break
|
|
}
|
|
posts = append(posts, newPosts...)
|
|
}
|
|
|
|
return posts, nil
|
|
}
|
|
|
|
// GetAllPosts retrieves all available posts using the provided post builder.
|
|
func (c *Client) GetAllPosts(postBuilder builder.PostsBuilder) ([]model.Post, error) {
|
|
// Retrieves all available posts using the provided post builder.
|
|
return c.GetNPosts(math.MaxInt, postBuilder)
|
|
}
|
|
|
|
func (c *Client) GetLatestPoolsDBExportDataAsBytes() (string, []byte, error) {
|
|
dbExportFileNameList, err := builder.NewGetDBExportListBuilder(c.RequestContext).Execute()
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
getDBExportFile := builder.NewGetDBExportFileBuilder(c.RequestContext)
|
|
|
|
filter := func(s string) bool { return strings.HasPrefix(s, "pools") }
|
|
filteredFileNameList := utils.SliceFilter(dbExportFileNameList, filter)
|
|
|
|
regex, err := regexp.Compile("\\d{4}\\-(0?[1-9]|1[012])\\-(0?[1-9]|[12][0-9]|3[01])*")
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
currentDate := time.Now()
|
|
|
|
duration := math.MaxFloat64
|
|
var fileName string
|
|
for _, listFileName := range filteredFileNameList {
|
|
if !regex.MatchString(listFileName) {
|
|
return "", nil, nil
|
|
}
|
|
|
|
fileDate, err := time.Parse("2006-04-02", regex.FindString(listFileName))
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
if currentDate.Sub(fileDate).Seconds() < duration {
|
|
duration = currentDate.Sub(fileDate).Seconds()
|
|
fileName = listFileName
|
|
}
|
|
}
|
|
|
|
rawFile, err := getDBExportFile.SetFile(fileName).Execute()
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
return fileName, rawFile, nil
|
|
}
|
|
|
|
func (c *Client) GetLatestPoolsDBExportDataAsStruct() ([]*model.Pool, error) {
|
|
var pools []*model.Pool
|
|
|
|
_, data, err := c.GetLatestPoolsDBExportDataAsBytes()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
zipReader, err := gzip.NewReader(bytes.NewReader(data))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
defer zipReader.Close()
|
|
|
|
// Create a CSV reader
|
|
reader := csv.NewReader(zipReader)
|
|
|
|
err = gocsv.UnmarshalCSV(reader, &pools)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return pools, nil
|
|
}
|
|
|
|
func (c *Client) GetLatestPostsDBExportDataAsBytes() (string, []byte, error) {
|
|
log.Println("Please wait while the download is in progress for the post export... (file over 1GB) ")
|
|
dbExportFileNameList, err := builder.NewGetDBExportListBuilder(c.RequestContext).Execute()
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
getDBExportFile := builder.NewGetDBExportFileBuilder(c.RequestContext)
|
|
|
|
filter := func(s string) bool { return strings.HasPrefix(s, "posts") }
|
|
filteredFileNameList := utils.SliceFilter(dbExportFileNameList, filter)
|
|
|
|
regex, err := regexp.Compile("\\d{4}\\-(0?[1-9]|1[012])\\-(0?[1-9]|[12][0-9]|3[01])*")
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
currentDate, err := time.Parse("2006-04-02", time.Now().Format("2006-01-02"))
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
duration := math.MaxFloat64
|
|
var fileName string
|
|
for _, listFileName := range filteredFileNameList {
|
|
if !regex.MatchString(listFileName) {
|
|
return "", nil, nil
|
|
}
|
|
|
|
fileDate, err := time.Parse("2006-04-02", regex.FindString(listFileName))
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
if currentDate.Sub(fileDate).Seconds() < duration {
|
|
duration = currentDate.Sub(fileDate).Seconds()
|
|
fileName = listFileName
|
|
}
|
|
}
|
|
|
|
rawFile, err := getDBExportFile.SetFile(fileName).Execute()
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
|
|
return fileName, rawFile, nil
|
|
}
|
|
|
|
func (c *Client) GetLatestPostsDBExportDataAsStruct() ([]*model.Post, error) {
|
|
var post []*model.Post
|
|
|
|
// TODO: Implement this function; tags in CSV are just a string with no categories assignment, needs special parsing
|
|
|
|
return post, nil
|
|
}
|