PrivacyScanData/archive/new2/bulk-downloader-for-reddit-development/bdfr/oauth2.py

#!/usr/bin/env python3
# coding=utf-8

import configparser
import logging
import random
import re
import socket
from pathlib import Path

import praw
import requests

from bdfr.exceptions import BulkDownloaderException, RedditAuthenticationError

logger = logging.getLogger(__name__)


class OAuth2Authenticator:

    def __init__(self, wanted_scopes: set[str], client_id: str, client_secret: str):
        self._check_scopes(wanted_scopes)
        self.scopes = wanted_scopes
        self.client_id = client_id
        self.client_secret = client_secret

    @staticmethod
    def _check_scopes(wanted_scopes: set[str]):
        response = requests.get('https://www.reddit.com/api/v1/scopes.json',
                                headers={'User-Agent': 'fetch-scopes test'})
        known_scopes = [scope for scope, data in response.json().items()]
        known_scopes.append('*')
        for scope in wanted_scopes:
            if scope not in known_scopes:
                raise BulkDownloaderException(f'Scope {scope} is not known to reddit')

    @staticmethod
    def split_scopes(scopes: str) -> set[str]:
        scopes = re.split(r'[,: ]+', scopes)
        return set(scopes)

    def retrieve_new_token(self) -> str:
        reddit = praw.Reddit(
            redirect_uri='http://localhost:7634',
            user_agent='obtain_refresh_token for BDFR',
            client_id=self.client_id,
            client_secret=self.client_secret)
        state = str(random.randint(0, 65000))
        url = reddit.auth.url(self.scopes, state, 'permanent')
        logger.warning('Authentication action required before the program can proceed')
        logger.warning(f'Authenticate at {url}')

        client = self.receive_connection()
        data = client.recv(1024).decode('utf-8')
        param_tokens = data.split(' ', 2)[1].split('?', 1)[1].split('&')
        params = {key: value for (key, value) in [token.split('=') for token in param_tokens]}

        if state != params['state']:
            self.send_message(client)
            raise RedditAuthenticationError(f'State mismatch in OAuth2. Expected: {state} Received: {params["state"]}')
        elif 'error' in params:
            self.send_message(client)
            raise RedditAuthenticationError(f'Error in OAuth2: {params["error"]}')

        self.send_message(client, "<script>alert('You can go back to terminal window now.')</script>")
        refresh_token = reddit.auth.authorize(params["code"])
        return refresh_token

    @staticmethod
    def receive_connection() -> socket.socket:
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(('0.0.0.0', 7634))
        logger.log(9, 'Server listening on 0.0.0.0:7634')

        server.listen(1)
        client = server.accept()[0]
        server.close()
        logger.log(9, 'Server closed')

        return client

    @staticmethod
    def send_message(client: socket.socket, message: str = ''):
        client.send(f'HTTP/1.1 200 OK\r\n\r\n{message}'.encode('utf-8'))
        client.close()


class OAuth2TokenManager(praw.reddit.BaseTokenManager):
    def __init__(self, config: configparser.ConfigParser, config_location: Path):
        super(OAuth2TokenManager, self).__init__()
        self.config = config
        self.config_location = config_location

    def pre_refresh_callback(self, authorizer: praw.reddit.Authorizer):
        if authorizer.refresh_token is None:
            if self.config.has_option('DEFAULT', 'user_token'):
                authorizer.refresh_token = self.config.get('DEFAULT', 'user_token')
                logger.log(9, 'Loaded OAuth2 token for authoriser')
            else:
                raise RedditAuthenticationError('No auth token loaded in configuration')

    def post_refresh_callback(self, authorizer: praw.reddit.Authorizer):
        self.config.set('DEFAULT', 'user_token', authorizer.refresh_token)
        with open(self.config_location, 'w') as file:
            self.config.write(file, True)
        logger.log(9, f'Written OAuth2 token from authoriser to {self.config_location}')