示例#1
0
class AppConnect:
    """App connection object.

    A wrapper for requests BaseUrlSession to hold Atlassian keys across command runs.

    Parameters
    ----------
    server:          base url of app server.
    username:        username for connection.
    password:        password for connection.
    cookie_store:    path to file for cookie_store.
    session_headers: default headers added to every call.
    """
    _server: str
    username: str
    _password: str
    session: BaseUrlSession = None
    auth: HTTPBasicAuth = None
    _response: requests = None
    cookie_store: os.path = None

    def __init__(self,
                 server: str,
                 username: str = None,
                 password: str = None,
                 cookie_store: os.path = None,
                 session_headers: dict = None) -> None:
        self.server = server
        self.session = BaseUrlSession(base_url=server)

        if username:
            self.username = username

        if password:
            self.password = password

        if cookie_store:
            self.cookie_store = cookie_store

        if username and password:
            self.auth = HTTPBasicAuth(self.username, self.password)

        if session_headers:
            self.session.headers.update(session_headers)

        self.reload_cookies()

    @property
    def server(self):
        """server baseUrl for connection"""
        return self._server

    @server.setter
    def server(self, server: str):
        self._server = server
        if self.session:
            self.session.base_url = server

    @property
    def password(self):
        """password for connection."""
        return base64.decodebytes(self._password).decode()

    @password.setter
    def password(self, password: str):
        self._password = base64.encodebytes(password.encode())

    def get(self,
            api,
            headers: dict = None,
            params: dict = None,
            data: dict = None,
            auth: bool = False,
            allow_redirects=True):
        """send http get request.

        Parameters
        ----------
        api:        str url path appended to baseUrl.
        headers:    dict of headers.
        params:     dict of url query parameters.
        data:       dict of data to send.
        auth:       bool(False) send BasicAuth.
        allow_redirects

        Returns
        -------

        """
        # url = urljoin(self.server, api)
        url = api

        try:
            self._response = self.session.get(url,
                                              headers=headers,
                                              params=params,
                                              data=data,
                                              auth=self.auth if auth else None,
                                              allow_redirects=allow_redirects)
            self._response.raise_for_status()
        except requests.exceptions.ConnectionError as err:
            raise SystemExit(err)
        except requests.exceptions.Timeout as err:
            raise SystemExit(err)
        except requests.exceptions.TooManyRedirects as err:
            raise SystemExit(err)
        except requests.exceptions.HTTPError as err:
            raise SystemExit(err)

        return self.json_response(self._response)

    def delete(self,
               api,
               headers: dict = None,
               params=None,
               auth: bool = False):
        """send http delete request.

        Parameters
        ----------
        api:        str url path appended to baseUrl.
        headers:    dict of headers.
        params:     dict of url query parameters.
        auth:       bool(False) send BasicAuth.

        Returns
        -------
        ->json
        """
        url = api

        try:
            self._response = self.session.delete(
                url,
                headers=headers,
                params=params,
                auth=self.auth if auth else None)
            self._response.raise_for_status()
        except requests.exceptions.ConnectionError as err:
            raise SystemExit(err)
        except requests.exceptions.Timeout as err:
            raise SystemExit(err)
        except requests.exceptions.TooManyRedirects as err:
            raise SystemExit(err)
        except requests.exceptions.HTTPError as err:
            raise SystemExit(err)

        return self.json_response(self._response)

    def post(self,
             api: str,
             headers: dict = None,
             params: dict = None,
             data: dict = None,
             auth: bool = False,
             allow_redirects: bool = True):
        """send http post request.

        Parameters
        ----------
        api:        str url path appended to baseUrl.
        headers:    dict of headers.
        params:     dict of url query parameters.
        data:       dict of data to send.
        auth:       bool(False) send BasicAuth.
        allow_redirects

        Returns
        -------
        ->json
        """
        # url = urljoin(self.server, api)
        url = api

        try:
            self._response = self.session.post(
                url,
                headers=headers,
                params=params,
                data=data,
                auth=self.auth if auth else None,
                allow_redirects=allow_redirects)
            # self._response.raise_for_status()
        except requests.exceptions.ConnectionError as err:
            raise SystemExit(err)
        except requests.exceptions.Timeout as err:
            raise SystemExit(err)
        except requests.exceptions.TooManyRedirects as err:
            raise SystemExit(err)
        # except requests.exceptions.HTTPError as err:
        #     raise SystemExit(err)

        return self.json_response(self._response)

    def put(self,
            api: str,
            headers: dict = None,
            params: dict = None,
            data: dict = None,
            auth: bool = False):
        """send http put request.

        Parameters
        ----------
        api:        str url path appended to baseUrl.
        headers:    dict of headers.
        params:     dict of url query parameters.
        data:       dict of data to send.
        auth:       bool(False) send BasicAuth.

        Returns
        -------
        ->json
        """
        url = api

        try:
            self._response = self.session.put(url,
                                              headers=headers,
                                              params=params,
                                              data=data,
                                              auth=self.auth if auth else None)
            self._response.raise_for_status()
        except requests.exceptions.ConnectionError as err:
            raise SystemExit(err)
        except requests.exceptions.Timeout as err:
            raise SystemExit(err)
        except requests.exceptions.TooManyRedirects as err:
            raise SystemExit(err)
        except requests.exceptions.HTTPError as err:
            raise SystemExit(err)

        return self.json_response(self._response)

    def json_response(self, res: requests):
        """Always return a json response.

        Parameters
        ----------
        res:    requests response.

        Returns
        -------
        ->json
        """
        _json = None

        if res.ok:
            if res.cookies:
                self.session.cookies.update(res.cookies)
            self.cache_cookies()

            try:
                _json = res.json()
            except JSONDecodeError as err:
                SystemExit(err)

        if not _json:
            if res.ok:
                _json = json.dumps({
                    'success':
                    self._response.ok,
                    'status code':
                    self._response.status_code,
                    'elapsed seconds':
                    self._response.elapsed.seconds
                })
            else:
                _json = json.dumps({
                    'ok': self._response.ok,
                    'status_code': self._response.status_code,
                    'reason': self._response.text,
                    'request-url': self._response.request.url,
                    'request-method': self._response.request.method,
                    'text': self._response.text,
                    'redirect': self._response.is_redirect,
                    'elapsed': self._response.elapsed.seconds
                })

        return _json

    def update_cookies(self, cookies: dict = None):
        """add cookie(s) to cookie jar.

        Parameters
        ----------
        cookies
        """
        self.session.cookies.update(cookies)
        self.cache_cookies()

    def cache_cookies(self):
        """cache cookies to file."""
        if self.session.cookies:
            with open(self.cookie_store, 'wb') as f:
                pickle.dump(self.session.cookies, f)

    def reload_cookies(self):
        """reload cookies from file."""
        if os.path.isfile(self.cookie_store):
            with open(self.cookie_store, 'rb') as f:
                self.session.cookies.update(pickle.load(f))
示例#2
0
class RCTFAdminV1:

    session: requests.Session

    def __init__(self, endpoint: str, login_token: Optional[str]):
        self.session = BaseUrlSession(urljoin(endpoint, "api/v1/admin/"))

        if login_token is not None:
            login_resp = requests.post(
                urljoin(endpoint, "api/v1/auth/login"), json={"teamToken": login_token}
            ).json()
            if login_resp["kind"] == "goodLogin":
                auth_token = login_resp["data"]["authToken"]
                self.session.headers["Authorization"] = f"Bearer {auth_token}"
            else:
                raise ValueError(
                    f"Invalid login_token provided (reason: {login_resp['kind']})"
                )

    @staticmethod
    def assertResponseKind(response: Any, kind: str) -> None:
        if response["kind"] != kind:
            raise RuntimeError(f"Server error: {response['kind']}")

    def list_challenges(self) -> List[Dict[str, Any]]:
        r = self.session.get("challs").json()
        self.assertResponseKind(r, "goodChallenges")
        return r["data"]

    def put_challenge(self, chall_id: str, data: Dict[str, Any]) -> None:
        r = self.session.put("challs/" + quote(chall_id), json={"data": data}).json()
        self.assertResponseKind(r, "goodChallengeUpdate")

    def delete_challenge(self, chall_id: str) -> None:
        r = self.session.delete("challs/" + quote(chall_id)).json()
        self.assertResponseKind(r, "goodChallengeDelete")

    def create_upload(self, uploads: Dict[str, bytes]) -> Dict[str, str]:
        """
        :param uploads: uploads {name: data}
        :return: urls {name: url}
        """
        if len(uploads) == 0:
            return {}
        payload = [
            {"name": name, "data": "data:;base64," + b64encode(data).decode()}
            for name, data in uploads.items()
        ]
        r = self.session.post("upload", json={"files": payload}).json()
        self.assertResponseKind(r, "goodFilesUpload")
        return {f["name"]: f["url"] for f in r["data"]}

    def get_url_for_files(self, files: Dict[str, str]) -> Dict[str, Optional[str]]:
        """
        :param files: files to get {name: sha256}
        :return: urls {name: url}
        """
        payload = [{"name": name, "sha256": sha256} for name, sha256 in files.items()]
        r = self.session.post("upload/query", json={"uploads": payload}).json()
        self.assertResponseKind(r, "goodUploadsQuery")
        return {f["name"]: f["url"] for f in r["data"]}