106 lines
4.0 KiB
Python

#!/usr/bin/env python3
"""Define the user class to populate and setup the download environment"""
import re
from random import shuffle
from pathlib import Path
from functions import load_config_variables
from functions import validate_x_link
from functions import parse_link
from functions import clean_cache
from functions import LOG
class User:
"""Populate the directory for each user"""
# pylint: disable=too-many-instance-attributes
def __init__(self, index) -> None:
config = load_config_variables()
self.config = config["users"][index] | config["global"]
self.name = self.config["name"]
self.sleep = self.config["sleep"]
# Directories
self.directories = {
str(key).replace("-dir", ""): Path(self.config[f"{key}"])
for key in filter(lambda x: re.search("-dir", x), self.config.keys())
}
self.directories["cache"] = self.directories["cache"] / self.name
self.directories["lists"] = self.directories["lists"] / self.name
# Files
self.dbs = {
"gallery": self.directories["databases"] / f"{self.name}.sqlite3",
"media": self.directories["databases"] / f"{self.name}_ytdl.txt",
}
# Lists
self.lists = {
"master": self.directories["lists"] / "watch.txt",
"push": self.directories["lists"] / "instant.txt",
"instagram": self.directories["cache"] / "instagram.txt",
"kemono": self.directories["cache"] / "kemono.txt",
"main": self.directories["cache"] / "main.txt",
}
def _create_directories(self) -> None:
"""Create user directories if they don't exist"""
clean_cache(self.directories["cache"])
# Create directories
for directory in self.directories.keys():
self.directories[directory].mkdir(parents=True, exist_ok=True)
# Check for the existence of core files
if not self.directories["lists"].is_dir():
LOG.error("Lists directory for user %s doesn't exist", self.name)
# dbs stands for databases, the archives.
for db in filter(lambda x: not self.dbs[x].is_file(), self.dbs.keys()):
self.dbs[db].touch()
for lst in filter(lambda x: not self.lists[x].is_file(), ["master", "push"]):
self.lists[lst].touch()
def append_list(self, name: str, line: str) -> None:
"""Appends a line into the given list"""
with open(self.lists[name], "a+", encoding="utf-8") as a_file:
a_file.write(line + "\n")
def _append_cache_list(self, line) -> None:
"""Writes the input line into it's respective list,
depending on what website it belongs to."""
if re.search("x", line):
self.append_list("main", validate_x_link(line))
elif re.search(r"kemono\.party", line):
self.append_list("kemono", line)
elif re.search("instagram", line):
self.append_list("instagram", line)
else:
self.append_list("main", line)
def list_manager(self) -> None:
"""Manage all the user list and create sub-lists"""
self._create_directories() # Call the function to create necesary cache dirs
with open(self.lists["master"], "r", encoding="utf-8") as r_file:
master_content = list(map(lambda x: x.rstrip(), r_file))
# Create temporary list files segmented per scrapper
shuffle(master_content)
for line in master_content:
self._append_cache_list(line)
def save_link(self, link: str) -> None:
"""Checks the master list against a new link
if unmatched, appends it to the end of the list"""
with open(self.lists["master"], "r", encoding="utf-8") as r_file:
links = r_file.read().lower()
if parse_link(link).lower() in links:
LOG.info("Gallery repeated, not saving")
return
LOG.info("New gallery, saving")
self.append_list("master", parse_link(link))