CourseRace.py

import requests
requests.packages.urllib3.disable_warnings()

from prettytable import PrettyTable
from concurrent.futures import ThreadPoolExecutor
import time
import os
import json
import argparse
import logging
import threading
from collections import deque # For a thread-safe queue

sourceUrl = 'https://xk.bit.edu.cn/yjsxkapp/sys/xsxkappbit/xsxkCourse/choiceCourse.do?_='
sourceUrl_vpn = 'https://webvpn.bit.edu.cn/https/77726476706e69737468656265737421e8fc0f9e2e2426557a1dc7af96/yjsxkapp/sys/xsxkappbit/xsxkCourse/choiceCourse.do?vpn-12-o2-xk.bit.edu.cn&_='

infoPage = 'https://xk.bit.edu.cn/yjsxkapp/sys/xsxkappbit/xsxkHome/loadPublicInfo_course.do'
infoPage_vpn = 'https://webvpn.bit.edu.cn/https/77726476706e69737468656265737421e8fc0f9e2e2426557a1dc7af96/yjsxkapp/sys/xsxkappbit/xsxkHome/loadPublicInfo_course.do?vpn-12-o2-xk.bit.edu.cn'

OutPlanCoursePage = 'https://xk.bit.edu.cn/yjsxkapp/sys/xsxkappbit/xsxkCourse/loadGxkCourseInfo.do?_='
OutPlanCoursePage_vpn = 'https://webvpn.bit.edu.cn/https/77726476706e69737468656265737421e8fc0f9e2e2426557a1dc7af96/yjsxkapp/sys/xsxkappbit/xsxkCourse/loadGxkCourseInfo.do?vpn-12-o2-xk.bit.edu.cn&_='

InPlanCoursePage = 'https://xk.bit.edu.cn/yjsxkapp/sys/xsxkappbit/xsxkCourse/loadJhnCourseInfo.do?_='
InPlanCoursePage_vpn = 'https://webvpn.bit.edu.cn/https/77726476706e69737468656265737421e8fc0f9e2e2426557a1dc7af96/yjsxkapp/sys/xsxkappbit/xsxkCourse/loadJhnCourseInfo.do?vpn-12-o2-xk.bit.edu.cn&_='

OutPlanCoursePath = './OutPlanCourses.json'
InPlanCoursePath = './InPlanCourses.json'

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36',
    'Cookie': ''  # add your cookie here
}

# add class info here
# this is examples
# you can copy it and change bjdm to your course
juzhen_zgc01_data = {
    'bjdm': '20231-17-1700002-1688866107858',
    'lx': '0',
    'csrfToken': '23b21ddb67914b3e81ae61923fd164aa'
}

courseList = [
    # add class info struct here

]

# A queue to hold courses that need to be chosen
course_queue = deque()
# A lock to protect access to the courseList when updating csrfToken
csrf_token_lock = threading.Lock()

# Define MAX_CONCURRENT_CHOICES globally, but it will be set in __main__
MAX_CONCURRENT_CHOICES = None # Will be set after parsing arguments
choice_semaphore = None # Will be initialized in __main__


def printErr(string):
    print('\033[31m' + string + '\033[0m')


def printOK(string):
    print('\033[32m' + string + '\033[0m')


def setVPN():
    global sourceUrl, infoPage, InPlanCoursePage, OutPlanCoursePage
    sourceUrl = sourceUrl_vpn
    infoPage = infoPage_vpn
    InPlanCoursePage = InPlanCoursePage_vpn
    OutPlanCoursePage = OutPlanCoursePage_vpn


def is_valid_json(json_str):
    try:
        json.loads(json_str)
        return True
    except json.JSONDecodeError as e:
        printErr("[-] Fail to catch courses. ERROR:" + str(e))
        return False


def postData(reqCourseList, req_data):
    try:
        res = requests.post(url=reqCourseList, data=req_data, headers=headers, verify=False)
        res.raise_for_status()
        return res
    except requests.exceptions.HTTPError as errh:
        printErr("[-] Fail to catch courses. HTTP ERROR:" + str(errh))
    except requests.exceptions.ConnectionError as errc:
        printErr("[-] Fail to catch courses. Connection ERROR:" + str(errc))
    except requests.exceptions.Timeout as errt:
        printErr("[-] Fail to catch courses. Timeout ERROR:" + str(errt))
    except requests.exceptions.RequestException as err:
        printErr("[-] Fail to catch courses. Unknown ERROR:" + str(err))

    return None


def getCourseList():

    req_data = {
        'query_keyword':        '',
        'query_kkyx':           '',
        'query_sfct':           '',
        'query_sfym':           '',
        'fixedAutoSubmitBug':   '',
        'pageIndex':            1,
        'pageSize':             1000,
        'sortField':            '',
        'sortOrder':            '',
    }

    print('[*] Try to catch courses out of plan...')

    timestamp = int(round(time.time() * 1000))
    reqCourseList = OutPlanCoursePage + str(timestamp)

    res = postData(reqCourseList, req_data)
    if not res:
        exit(1)
    if not is_valid_json(res.text):
        exit(1)

    with open(OutPlanCoursePath, 'w', encoding='utf8') as f:
        f.write(res.text)
    print('[+] Success. Courses have been saved in ' + OutPlanCoursePath)

    print('[*] Try to catch courses in plan...')

    timestamp = int(round(time.time() * 1000))
    reqCourseList = InPlanCoursePage + str(timestamp)

    res = postData(reqCourseList, req_data)
    if not res:
        exit(1)
    if not is_valid_json(res.text):
        exit(1)

    with open(InPlanCoursePath, 'w', encoding='utf8') as f:
        f.write(res.text)
    print('[+] Success. Courses have been saved in ' + InPlanCoursePath)


def findCourse(idList: list):
    with open(InPlanCoursePath, "r", encoding="utf8") as f:
        InPlanCourseInfo = f.read()
    InPlanCourseInfo = json.loads(InPlanCourseInfo)
    with open(OutPlanCoursePath, "r", encoding="utf8") as f:
        OutPlanCourseInfo = f.read()
    OutPlanCourseInfo = json.loads(OutPlanCourseInfo)

    targetList = []
    for id in idList:
        print("[*] Looking for course id:", id, "...")
        for info in InPlanCourseInfo['datas']:
            if id == info["KCDM"] and info["XQMC"] != "良乡校区" and ("非全" not in info["BJMC"]):
                targetList.append([info["KCMC"], info["RKJS"], "{}/{}".format(info["DQRS"], info["KXRS"])])
                courseList.append({'bjdm': info["BJDM"], 'lx': '0', 'csrfToken': ''})
        for info in OutPlanCourseInfo['datas']:
            if id == info["KCDM"] and info["XQMC"] != "良乡校区" and ("非全" not in info["BJMC"]):
                targetList.append([info["KCMC"], info["RKJS"], "{}/{}".format(info["DQRS"], info["KXRS"])])
                courseList.append({'bjdm': info["BJDM"], 'lx': '1', 'csrfToken': ''})

    if len(targetList) == 0:
        print("[!] No course found according to course id.")
        exit(0)

    table = PrettyTable()
    table.field_names = ['Name', 'Teachers', 'Chosen']
    table.add_rows(targetList)
    print("[+] Target courses shown as follow:")
    print(table)


def chooseCourse(course):
    global choice_semaphore # Declare that we're using the global semaphore

    # Acquire a semaphore slot before proceeding
    choice_semaphore.acquire()
    try:
        while True:
            timestamp = int(round(time.time() * 1000))
            courseUrl = sourceUrl + str(timestamp)
            res = requests.post(url=courseUrl, data=course, headers=headers, verify=False)
            res = json.loads(res.text)
            if(res["code"] == 1):
                printOK(f"[+] Course {course['bjdm']} is chosen! You can see on Web Browser!")
                # If a course is successfully chosen, we can stop trying for this course.
                break
            else:
                logging.debug(f"[{course['bjdm']}] {res['msg']}")
            time.sleep(0.01) # Small delay to prevent hammering the server
    finally:
        # Release the semaphore slot when done, regardless of success or failure
        choice_semaphore.release()


def start():
    print("[*] Start race...Please wait for several hours...")
    global MAX_CONCURRENT_CHOICES, choice_semaphore # Access global variables

    # Initialize the queue with all courses that need to be chosen
    for course in courseList:
        course_queue.append(course)

    # Use ThreadPoolExecutor to manage a fixed number of workers
    # The max_workers for the ThreadPoolExecutor should be based on MAX_CONCURRENT_CHOICES
    pool = ThreadPoolExecutor(max_workers=MAX_CONCURRENT_CHOICES)
    print(f"[*] Thread pool initialized with {MAX_CONCURRENT_CHOICES} workers.")

    # Submit all courses to the pool initially. The semaphore in chooseCourse will
    # ensure only MAX_CONCURRENT_CHOICES are actively making requests at any given time.
    futures = [pool.submit(chooseCourse, course) for course in course_queue]

    # This loop will periodically fetch the csrfToken and update the courses
    while True:
        try:
            res = requests.get(url=infoPage, headers=headers, verify=False)
            res.raise_for_status()
            csrfToken = json.loads(res.text)['csrfToken']

            # Use a lock when modifying shared `courseList` resources to prevent race conditions
            with csrf_token_lock:
                for course in courseList: # Update all courses in the original list
                    course['csrfToken'] = csrfToken
            logging.debug(f"[*] Updated csrfToken: {csrfToken}")

        except requests.exceptions.RequestException as e:
            printErr(f"[-] Failed to fetch CSRF token: {e}")
        except json.JSONDecodeError:
            printErr("[-] Failed to parse CSRF token response.")

        # Check if all courses have been chosen (futures are done)
        if all(f.done() for f in futures):
            printOK("[+] All target courses have been processed or chosen!")
            break

        time.sleep(60) # Wait for 60 seconds before fetching a new CSRF token

    pool.shutdown(wait=True) # Ensure all threads are finished


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='BIT Course Race. A script to help masters get courses.')
    parser.add_argument("-c", "--cookie",
                         type=str,
                         dest="cookie",
                         help="Cookie copied from your web browser(after logging in sucessfully)")
    parser.add_argument("-i", "--courseID",
                         type=str,
                         dest="courseID",
                         nargs='+',
                         help="ID of courses, split with space")
    parser.add_argument("-v", "--vpn",
                         dest="vpn",
                         action='store_true',
                         help="if you choose course through webvpn, then use this")
    parser.add_argument("-d", "--debug",
                         dest="debug",
                         action='store_true',
                         help="if you want to show debug messages, then use this")
    parser.add_argument("-t", "--threads",
                         type=int,
                         dest="threads",
                         default=None, # Default to None, so we can detect if it was provided
                         help="Maximum number of threads to use (will not exceed hardware threads)")

    args = parser.parse_args()
    headers['Cookie'] = args.cookie

    if args.vpn is True:
        setVPN()

    if args.debug is True:
        logging.basicConfig(level = logging.DEBUG,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    else:
        logging.basicConfig(format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    # Determine the maximum number of threads
    hardware_threads = os.cpu_count() if os.cpu_count() else 1 # Fallback to 1 if os.cpu_count() is None

    if args.threads is not None and args.threads > 0:
        if args.threads > hardware_threads:
            print(f"[*] Warning: User-specified threads ({args.threads}) exceed hardware threads ({hardware_threads}). Using {hardware_threads} threads.")
            MAX_CONCURRENT_CHOICES = hardware_threads
        else:
            MAX_CONCURRENT_CHOICES = args.threads
            print(f"[*] Using user-specified {args.threads} threads.")
    else:
        MAX_CONCURRENT_CHOICES = hardware_threads
        print(f"[*] No valid thread count specified. Using hardware threads: {hardware_threads}.")

    # Initialize the semaphore after MAX_CONCURRENT_CHOICES is determined
    choice_semaphore = threading.Semaphore(MAX_CONCURRENT_CHOICES)

    getCourseList()

    findCourse(args.courseID)

    start()
    
    # res = requests.get(url=infoPage, headers=headers, verify=False)
    # csrfToken = json.loads(res.text)['csrfToken']
    # for course in courseList:
    #     course['csrfToken'] = csrfToken