aliyun-share

#! /bin/env python3

# Requirements:
# 1. [aligo](https://github.com/foyoux/aligo), which is available on pypi.org
# 2. [aria2c](https://aria2.github.io), optional dependency for downloading
# 2. An (aliyun)[https://www.aliyundrive.com/] account

# Usage Example:
# 1. Store share link index data into PostgreSQL within category book
# > python3 ./aliyun-share https://www.aliyundrive.com/s/wnfVTd1u92T book
# 2. Invoke researches from database within category book
# > python3 ./aliyun-share 'YOUR SEARCHES' book

import re
import sys
import tempfile
from type import CATEGORIES
from aligo import Aligo
from psycopg2 import connect
from psycopg2.sql import SQL, Identifier
from download import Download
from datetime import datetime
import logging

# Call this script without arguments to login or persist login info
if len(sys.argv) == 1:
    # refresh login token
    from aligo import Auth
    Auth(name='aligo')._refresh_token()
    sys.exit()

# Use global variables to refresh share_token after around 2 hours
update_interval = 7100
expired_time = datetime.now().timestamp() + update_interval


conn = connect("dbname=alibrary")
cur = conn.cursor()

category = CATEGORIES[0]

if len(sys.argv) == 3 and sys.argv[2] in CATEGORIES:
    category = sys.argv[2]

print("Current category is", category)

cur.execute(SQL("CREATE TABLE IF NOT EXISTS {} (id char(52) PRIMARY KEY, name TEXT, size INTEGER, tsv TSVECTOR GENERATED ALWAYS AS (to_tsvector('jiebacfg', name)) STORED);").format(Identifier(category)))
cur.execute(SQL("CREATE INDEX IF NOT EXISTS {} ON {} USING GIN (tsv);").format(Identifier(category + '_idx'), Identifier(category)))
conn.commit()


def get_share_id(
        share_id_data_file: str = tempfile.gettempdir() + '/share_id') -> str:
    match_share_id = re.match(
        r"https://www.aliyundrive.com/s/([^/]+)(/folder)?", sys.argv[1])
    if match_share_id:
        share_id = match_share_id.group(1)
        f = open(share_id_data_file, 'w')
        f.write(share_id)
        f.close()
        return share_id


def get_file_id(share_id: str = None) -> str:
    match_file_id = re.match(r".*([0-9a-f]{40})$", sys.argv[1])
    if match_file_id:
        return match_file_id.group(1)
    else:
        return ali.get_share_info(share_id=share_id).file_infos[0].file_id


def record(share_info, share_token: str = None):

    if share_token is None:
        share_token: str = str(ali.get_share_token(share_info.share_id))

    global expired_time
    if datetime.now().timestamp() > expired_time:
        expired_time = datetime.now().timestamp() + update_interval
        record(share_info)
        return

    if share_info.type == 'folder':
        cur.execute("SELECT EXISTS(SELECT 1 FROM folder WHERE id = %s) LIMIT 1;", (share_info.file_id,))
        if cur.fetchone()[0]:
            print("Folder", share_info.name, "already in the database")
            return
        file_list = ali.get_share_file_list(share_id=share_info.share_id,
                                            share_token=share_token,
                                            parent_file_id=share_info.file_id)
        for file in file_list:
            try:
                record(file, share_token=share_token)
            except AttributeError:
                print(file)
        cur.execute("INSERT INTO folder (id) VALUES (%s);", (share_info.file_id,))
        print('Add folder ', share_info.name)
        conn.commit()
    elif share_info.type == 'file':
        id = share_info.share_id + ':' + share_info.file_id
        cur.execute(SQL("INSERT INTO {} (id, name, size) VALUES (%s, %s, %s) ON CONFLICT DO NOTHING;").format(Identifier(category)),
                    (id, share_info.name, share_info.size // 1024))


share_id = get_share_id()
if share_id:
    ali = Aligo(level=logging.ERROR)
    share_token = ali.get_share_token(share_id)
    file_id = get_file_id(share_id)
    share_info = ali.get_share_file(share_id=share_id,
                                    file_id=file_id,
                                    share_token=str(share_token))
    record(share_info)
    conn.commit()
    cur.close()
    conn.close()
    if share_info.type == 'file':
        Download(share_info=share_info)
else:
    from prompt import SearchPrompt
    SearchPrompt(search_text=sys.argv[1], category=category)