149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
import asyncio
|
|
import re
|
|
import sys
|
|
from typing import Union
|
|
|
|
import aiohttp
|
|
from bs4 import BeautifulSoup
|
|
|
|
from .helpers import SB, random_string
|
|
from .types import Language
|
|
|
|
|
|
class CodeforcesClient:
|
|
def __init__(self, cookies: dict[str, str], csrf_token: str, ftaa: str, user_agent: str):
|
|
self.cookies = cookies
|
|
self.user_agent = user_agent
|
|
|
|
self._csrf_token = csrf_token
|
|
self._ftaa = ftaa
|
|
|
|
|
|
async def __aenter__(self):
|
|
self._session = aiohttp.ClientSession()
|
|
self._session.cookie_jar.update_cookies(self.cookies)
|
|
self._session.headers['User-Agent'] = self.user_agent
|
|
return self
|
|
|
|
|
|
async def __aexit__(self, *_):
|
|
await self.close()
|
|
|
|
|
|
@staticmethod
|
|
def auth(username: str, password: str): # TODO: error checking
|
|
with SB() as sb:
|
|
sb.uc_open_with_reconnect("https://codeforces.com/enter", 2)
|
|
sb.uc_gui_click_captcha()
|
|
|
|
sb.click('a[href="/enter?back=%2F"]')
|
|
|
|
sb.type("#handleOrEmail", username)
|
|
sb.type("#password", password)
|
|
|
|
sb.click("#remember")
|
|
sb.click('input[type="submit"]')
|
|
|
|
sb.wait_for_element("#jGrowl")
|
|
|
|
cookies = sb.get_cookies()
|
|
user_agent = sb.get_user_agent()
|
|
source = sb.get_page_source()
|
|
|
|
match = re.findall(r'<meta name="X-Csrf-Token" content="([a-f0-9]+)">', source)
|
|
assert match
|
|
csrf = match[0]
|
|
|
|
match = re.findall(r'window._ftaa = "(\w+)";', source)
|
|
assert match
|
|
ftaa = match[0]
|
|
|
|
cookies = {cookie['name']: cookie['value'] for cookie in cookies}
|
|
|
|
return CodeforcesClient(cookies, csrf, ftaa, user_agent)
|
|
|
|
|
|
@staticmethod
|
|
def find_error(body: str) -> Union[str, None]:
|
|
r = re.compile(r'error[a-zA-Z_\-\ ]*">(.*?)</span>')
|
|
m = re.findall(r, body)
|
|
return m[0] if m else None
|
|
|
|
|
|
async def submit(self, code: str, contest: int, task: str, lang: Language): # TODO: error checking
|
|
r = await self._session.post(f"https://codeforces.com/contest/{contest}/submit?csrf_token={self._csrf_token}", data = {
|
|
"csrf_token": self._csrf_token,
|
|
"ftaa": self._ftaa,
|
|
"bfaa": random_string(32, '[a-f0-9]'),
|
|
"action": "submitSolutionFormSubmitted",
|
|
"submittedProblemIndex": task, # TODO: idx saving
|
|
"programTypeId": lang.value,
|
|
"source": code,
|
|
"tabSize": "4",
|
|
"sourceFile": "",
|
|
"_tta": "239",
|
|
})
|
|
|
|
err = self.find_error(await r.text())
|
|
|
|
if err:
|
|
print(err, file=sys.stderr)
|
|
|
|
|
|
async def parse_tests(self, contest: int, tasks: list[str]):
|
|
futures = [self._session.get(f"https://codeforces.com/contest/{contest}/problem/{task}") for task in tasks]
|
|
requests = await asyncio.gather(*futures)
|
|
|
|
for idx, r in enumerate(requests):
|
|
soup = BeautifulSoup(await r.text(), 'html.parser')
|
|
test = soup.select_one('div.sample-test')
|
|
assert test
|
|
|
|
# funny line haha
|
|
inputs = ['\n'.join([el.text for el in input.select('div.test-example-line')] or [input.text]).strip() for input in test.select('div.input > pre')]
|
|
outputs = ['\n'.join([el.text for el in output.select('div.test-example-line')] or [output.text]).strip() for output in test.select('div.output > pre')]
|
|
|
|
yield {tasks[idx].lower(): [*zip(inputs, outputs)]}
|
|
|
|
|
|
async def parse_tasks(self, contest: int):
|
|
r = await self._session.get(f"https://codeforces.com/contest/{contest}")
|
|
soup = BeautifulSoup(await r.text(), 'html.parser')
|
|
|
|
table = soup.select_one('table.problems')
|
|
assert table
|
|
|
|
elements = table.select('tr')[1:]
|
|
|
|
out = []
|
|
|
|
for item in elements:
|
|
status = None
|
|
if 'class' in item.attrs:
|
|
status = item.attrs['class'][0]
|
|
inner_items = item.select('td')
|
|
assert inner_items
|
|
name_and_constraints = inner_items[1].select('div > div')
|
|
id = inner_items[0].text.strip()
|
|
|
|
out.append({
|
|
'id': id,
|
|
'name': name_and_constraints[0].text.strip(),
|
|
'status': status,
|
|
'constraints': [*name_and_constraints[1]][2].text.strip(),
|
|
})
|
|
|
|
return out
|
|
|
|
|
|
def to_dict(self):
|
|
return {
|
|
'ftaa': self._ftaa,
|
|
'cookies': self.cookies,
|
|
'user_agent': self.user_agent,
|
|
'csrf_token': self._csrf_token
|
|
}
|
|
|
|
|
|
async def close(self):
|
|
await self._session.close()
|