from mycqu.auth import login
from requests import Session, Response

USERNAME = '07102028'
PASSWORD = 'Zhud626300'

def access_sso_service(session: Session, service: str):
    class ArriveService(Exception):
        def __init__(self, url: str):
            super().__init__()
            self.url: str = url

    def response_hook(response: Response, *args, **kwargs):
        #print(response.status_code, response.url, response.headers.get('Location'))
        if response.url.startswith("https://sso.cqu.edu.cn/login") and response.status_code != 302:
            assert '<div class="code">adapter</div>' in response.text
            response.headers['Location'] = './clientredirect?client_name=adapter'
            response.status_code = 302
            return response
        if response.url.startswith("https://sso.cqu.edu.cn/login") and response.status_code == 302 and response.headers['Location'].startswith(service):
            raise ArriveService(response.headers['Location'])

    try:
        assert not session.get("https://sso.cqu.edu.cn/login",
                               params={"service": service},
                               allow_redirects=True,
                               hooks={"response": response_hook}
                               )
    except ArriveService as arrive_service:
        return arrive_service.url

session = Session()
login(session, USERNAME, PASSWORD)
print(access_sso_service(session, "https://atrust.cqu.edu.cn:443/passport/v1/auth/cas"))
