niming_igapi/grpcServer/grpcServer.py

# from concurrent import futures

import grpc
from cachetools import cached, TTLCache

from ig import IG
from db import dbhelper
from utils.const import GRPC_ACCINFO_CACHE, GRPC_RELOGIN_LIMIT
from grpcServer import anoth
from grpcServer.protobuf import igapi_pb2_grpc
from grpcServer.protobuf.igapi_pb2 import Request, Reply

# call account info  / login
cache_accinfo = TTLCache(maxsize=1, ttl=GRPC_ACCINFO_CACHE)
@cached(cache_accinfo)
def call_IG_account_info():
    result = IG.account_info()
    return result

cache_login = TTLCache(maxsize=1, ttl=GRPC_RELOGIN_LIMIT)
@cached(cache_login)
def call_IG_login():
    result = IG.login()
    return result


# object
# 考慮一下如果同時發起多的請求，asyncio可能會搞到被ban號(IG)
class IGAPI_Server(igapi_pb2_grpc.IGAPIServicer):
    async def account_info(self, request: Request, context) -> Reply:
        print("[*] Request: account_info")
        account = call_IG_account_info()
        if account:
            result = {
                "username":account["username"],
                "full_name":account["full_name"],
                "email":account["email"]
            }
            return Reply(err=0, result=result)
        else:
            return Reply(err=1, result={"error":"IG.account_info returned None"})

    
    async def login(self, request: Request, context) -> Reply:
        print("[*] Request: login")
        if len(cache_login): # cache has not expired
            print("[*] Login: Cooldown")
            return Reply(err=1, result={"error":"Cooldown"})
        else:
            login = call_IG_login()
            if login:
                return Reply(err=0, result={"result":"Login Successed"})
            else:
                return Reply(err=1, result={"error":"Login Failed"})
    
    
    async def upload(self, request: Request, context) -> Reply:
        print("[*] Request: upload")
        aid = request.code
        
        # 檢查 - 可見
        article, code = dbhelper.solo_article_fetcher(role="general", key=aid) # visible -> post
        if code != 200:
            return Reply(err=1, result={"error":"Post not found"})
        
        # 檢查 - 已經在 Queue 內
        if anoth.task["upload-"+str(aid)]:
            return Reply(err=1, result={"error":"Request is already in queue"})
        
        # 檢查 - Queue 內有請求刪除同目標
        if anoth.task["delete-"+str(aid)]:
            anoth.task.pop("delete-"+str(aid))
            return Reply(err=0, result={"result":"Canceled delete post request"})
        
        # 檢查 - 已經上傳過
        if article["igid"]:
            return Reply(err=1, result={"error":"Already Posted"})
        
        # put into queue
        anoth.task["upload-"+str(aid)] = {"aid":aid}
        
        return Reply(err=0, result={"result":"Put into queue"})
    
    
    async def delete(self, request: Request, context) -> Reply:
        print("[*] Request: delete")
        # article id
        aid = request.code
        # igid from args
        if request.args:
            igid = request.args[0]
        else:
            return Reply(err=1, result={"error":"Invalid Arguments"})
        
        # 檢查 - 已經在 Queue 內
        if anoth.task["delete-"+str(aid)]:
            return Reply(err=1, result={"error":"Request is already in queue"})
        
        # 檢查 - Queue 內有請求上傳同目標
        if anoth.task["upload-"+str(aid)]:
            anoth.task.pop("upload-"+str(aid))
            return Reply(err=0, result={"result":"Canceled upload post request"})
        
        # put into queue
        anoth.task["delete-"+str(aid)] = {"aid":aid, "code":igid}
        
        return Reply(err=0, result={"result":"Put into queue"})
    
    
    async def queue(self, request:Request, context) -> Reply:
        print("[*] Request: queue")
        t = anoth.task.items()
        reply = { _[0]:str(_[1]["aid"]) for _ in t }
        return Reply(err=0, result=reply)
    
    
    async def setting(self, request:Request, context) -> Reply:
        # not done
        print("[*] Request: setting")
        return Reply(err=1, result={"error":"Not Done"})


# start server
async def serve() -> None:
    server = grpc.aio.server()
    igapi_pb2_grpc.add_IGAPIServicer_to_server(
        IGAPI_Server(), server
    )
    server.add_insecure_port("[::]:50051")
    await server.start()
    print("[*] gRPC Server listening on 0.0.0.0:50051")
    await server.wait_for_termination()