增加pat提交题目的代码

同时增加了获取csrf token的逻辑

main.py

@@ -1,9 +1,16 @@
 # coding=utf-8
+import time
 from robots.pat import PATRobot
+from robots.utils import LANGUAGE
 
 
 p = PATRobot()
-# p.cookies = p.login("virusdefender", "092122302asdf")
+p.login("virusdefender", "092122302asdf")
 # print(p.is_logged_in)
-for k, v in p.get_problem("http://www.patest.cn/contests/pat-t-practise/1005").items():
-    print("[[[", k, "]]]", v)
+# for i in range(10, 50):
+#     print(i)
+#     for k, v in p.get_problem("http://www.patest.cn/contests/pat-b-practise/10" + str(i)).items():
+#         print("[[[", k, "]]]", v)
+#     time.sleep(3)
+# p.get_token()
+print(p.submit("http://www.patest.cn/contests/32/1001/submit", language=LANGUAGE.C, code="123"))

robots/exceptions.py

@@ -8,4 +8,8 @@ class RequestFailed(Exception):
 
 
 class RegexError(Exception):
+    pass
+
+
+class SubmitProblemFailed(Exception):
     pass

robots/pat.py

@@ -1,15 +1,26 @@
 # coding=utf-8
 import re
+import urllib
 from .robot import Robot
-from .exceptions import AuthFailed, RequestFailed, RegexError
+from .exceptions import AuthFailed, RequestFailed, RegexError, SubmitProblemFailed
+from .utils import LANGUAGE
 
 
 class PATRobot(Robot):
+    def __init__(self):
+        super().__init__()
+        self.token = ""
+
     def check_url(self, url):
-        regex = "^http://www.patest.cn/contests/pat-(a|b|t)-practise/1\d{3}$"
+        regex = r"^http://www.patest.cn/contests/pat-(a|b|t)-practise/1\d{3}$"
         return re.compile(regex).match(url) is not None
 
+    def _get_token(self):
+        r = self.get("http://www.patest.cn/contests", cookies=self.cookies)
+        self.token = re.compile(r"<meta content=\"(.*)\" name=\"csrf-token\" />").findall(r.text)[0]
+
     def login(self, username, password):
+
         r = self.post("http://www.patest.cn/users/sign_in",
                       data={"utf8": "✓",
                             "user[handle]": username,
@@ -21,7 +32,9 @@ class PATRobot(Robot):
         # 登陆成功会重定向到首页,否则200返回错误页面
         if r.status_code != 302:
             raise AuthFailed()
-        return dict(r.cookies)
+
+        self.cookies = dict(r.cookies)
+        self._get_token()
 
     @property
     def is_logged_in(self):
@@ -33,7 +46,7 @@ class PATRobot(Robot):
     def get_problem(self, url):
         if not self.check_url(url):
             raise RequestFailed("Invalid PAT url")
-        problem_id = "pat-" + "-".join(re.compile("^http://www.patest.cn/contests/pat-(a|b|t)-practise/(\d{4})$").findall(url)[0])
+        problem_id = "pat-" + "-".join(re.compile(r"^http://www.patest.cn/contests/pat-(a|b|t)-practise/(\d{4})$").findall(url)[0])
         regex = {"title": r"<div id=\"body\" class=\"span-22 last\">\s*<h1>(.*)</h1>",
                  "time_limit": r"<div class='key'>\s*时间限制\s*</div>\s*<div class='value'>\s*(\d+) ms",
                  "memory_limit": r"<div class='key'>\s*内存限制\s*</div>\s*<div class='value'>\s*(\d+) kB",
@@ -41,7 +54,7 @@ class PATRobot(Robot):
                  "input_description": r"<b>\s*(?:Input|Input Specification:|输入格式：)\s*</b>([\s\S]*?)<b>\s*(?:Output|Output Specification:|输出格式：)\s*</b>",
                  "output_description": r"<b>\s*(?:Output|Output Specification:|输出格式：)\s*</b>([\s\S]*?)<b>\s*(?:Sample Input|输入样例).*</b>",
                  "samples": r"<b>\s*(?:Sample Input|输入样例)\s*(?P<t_id>\d?).?</b>\s*<pre>([\s\S]*?)</pre>\s+<b>(?:Sample Output|输出样例)\s?(?P=t_id).?</b>\s*<pre>([\s\S]*?)</pre>"}
-        data= self._regex_page(url, regex)
+        data = self._regex_page(url, regex)
         data["id"] = problem_id
         return data
 
@@ -63,3 +76,19 @@ class PATRobot(Robot):
                 data[k] = tmp
         return data
 
+    def submit(self, url, language, code):
+        if language == LANGUAGE.C:
+            language = "3"
+        elif language == LANGUAGE.CPP:
+            language = "2"
+        else:
+            language = "10"
+        r = self.post(url, data={"utf8": "✓", "compiler_id": language, "code": code,
+                                 "authenticity_token": self.token},
+                      cookies=self.cookies,
+                      headers={"Referer": "http://www.patest.cn/",
+                               "Content-Type": "application/x-www-form-urlencoded"})
+        if r.status_code != 302:
+            raise SubmitProblemFailed("Failed to submit problem, url: %s, status code: %d" % (url, r.status_code))
+        return str(re.compile("http://www.patest.cn/submissions/(\d+)").findall(r.headers["Location"])[0])
+

robots/robot.py

@@ -1,5 +1,6 @@
 # coding=utf-8
 import re
+import urllib
 import requests
 from .exceptions import RequestFailed, RegexError
 
@@ -27,6 +28,15 @@ class Robot(object):
         """
         raise NotImplementedError()
 
+    def submit(self, url, language, code):
+        """
+        :param url
+        :param language:
+        :param code:
+        :return: 提交id
+        """
+        raise NotImplementedError()
+
     def _request(self, method, url, **kwargs):
         kwargs["timeout"] = 10
 

robots/utils.py

@@ -0,0 +1,5 @@
+# coding=utf-8
+class LANGUAGE(object):
+    C = 0
+    CPP = 1
+    JAVA = 2