Blame for app.py

import csv
import hashlib
import io
import json
import os
import secrets
import threading
from pathlib import Path
from datetime import datetime, timezone

from authlib.integrations.starlette_client import OAuth
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi, hf_hub_download
from pydantic import BaseModel, Field
from starlette.middleware.sessions import SessionMiddleware

ROOT = Path(__file__).parent
DATA = Path(os.getenv('DATA_DIR', str(ROOT / 'data')))
DATA.mkdir(parents=True, exist_ok=True)
REPO = os.getenv('DATA_REPO')
TOKEN = os.getenv('DATA_TOKEN')
SECRET = os.getenv('SESSION_SECRET')
if not SECRET:
    if os.getenv('SPACE_HOST'):
        raise RuntimeError('Set the SESSION_SECRET Space secret before starting.')
    SECRET = secrets.token_hex(32)
HOST = os.getenv('SPACE_HOST')
BASE = f'https://{HOST}' if HOST else os.getenv('BASE_URL', 'http://localhost:7860')
SPACE_ID = os.getenv('SPACE_ID', 'Banaxi-Tech/HuggingPolls')
SHARE_BASE = f'https://huggingface.co/spaces/{SPACE_ID}' if HOST else BASE
LOCK = threading.RLock()
api = HfApi(token=TOKEN) if REPO else None
path = DATA / 'polls.json'
if REPO:
    from huggingface_hub.errors import EntryNotFoundError
    try:
        downloaded = hf_hub_download(REPO, 'polls.json', repo_type='dataset', token=TOKEN, force_download=True)
        path.write_bytes(Path(downloaded).read_bytes())
    except EntryNotFoundError:
        pass
state = json.loads(path.read_text()) if path.exists() else {'polls': {}}


def save(new):
    """Publish the full snapshot before acknowledging a write. Single worker only."""
    raw = json.dumps(new, ensure_ascii=False).encode()
    if api:
        try:
            api.upload_file(path_or_fileobj=raw, path_in_repo='polls.json', repo_id=REPO,
                            repo_type='dataset', commit_message='Save poll activity')
        except Exception:
            raise HTTPException(503, 'Storage is temporarily unavailable. Please try again.')
    temporary = path.with_suffix('.tmp')
    temporary.write_bytes(raw)
    temporary.replace(path)
    state.clear()
    state.update(new)


def copy_state():
    return json.loads(json.dumps(state))


def now():
    return datetime.now(timezone.utc).isoformat()


app = FastAPI()
app.add_middleware(SessionMiddleware, secret_key=SECRET, same_site='lax', https_only=bool(HOST), max_age=28800)
app.mount('/static', StaticFiles(directory=ROOT / 'static'), name='static')
oauth = OAuth()
if os.getenv('OAUTH_CLIENT_ID'):
    oauth.register('hf', client_id=os.environ['OAUTH_CLIENT_ID'],
                   client_secret=os.environ['OAUTH_CLIENT_SECRET'],
                   server_metadata_url=os.getenv('OPENID_PROVIDER_URL', 'https://huggingface.co').rstrip('/') + '/.well-known/openid-configuration',
                   client_kwargs={'scope': 'openid profile', 'code_challenge_method': 'S256'})


@app.middleware('http')
async def protect(request: Request, call_next):
    if request.method in {'POST', 'PATCH', 'DELETE'}:
        # JSON-only mutations plus a custom header prevent cross-site form submissions.
        if request.headers.get('x-huggingpolls') != '1':
            return Response('Missing request header', status_code=403)
        origin = request.headers.get('origin')
        if origin and origin.rstrip('/') != BASE.rstrip('/'):
            return Response('Invalid origin', status_code=403)
    response = await call_next(request)
    response.headers['X-Content-Type-Options'] = 'nosniff'
    response.headers['Referrer-Policy'] = 'same-origin'
    if request.url.path.startswith('/api/'):
        response.headers['Cache-Control'] = 'no-store'
    return response


def user(request):
    value = request.session.get('user')
    if not value:
        raise HTTPException(401, 'Sign in with Hugging Face to continue.')
    return value


def get_poll(poll_id):
    poll = state['polls'].get(poll_id)
    if not poll:
        raise HTTPException(404, 'This poll could not be found.')
    return poll


def owned(request, poll_id):
    identity = user(request)
    poll = get_poll(poll_id)
    if poll['owner_id'] != identity['id']:
        raise HTTPException(403, 'Only the poll creator can do this.')
    return poll


def voter_keys(request):
    guest = request.session.get('voter')
    if not guest:
        guest = secrets.token_urlsafe(32)
        request.session['voter'] = guest
    values = ['guest:' + guest]
    if request.session.get('user'):
        values.append('hf:' + request.session['user']['id'])
    return [hashlib.sha256((SECRET + value).encode()).hexdigest() for value in values]


@app.get('/')
def home():
    return FileResponse(ROOT / 'static/index.html')


@app.get('/api/me')
def me(request: Request):
    voter_keys(request)
    return {'user': request.session.get('user'), 'login_available': bool(os.getenv('OAUTH_CLIENT_ID'))}


@app.get('/auth/login')
async def login(request: Request, poll: str = ''):
    if not os.getenv('OAUTH_CLIENT_ID'):
        raise HTTPException(503, 'HF sign-in becomes available when deployed as a Space with hf_oauth enabled.')
    request.session['return_poll'] = poll if poll.isalnum() else ''
    return await oauth.hf.authorize_redirect(request, BASE + '/auth/callback')


@app.get('/auth/callback')
async def callback(request: Request):
    try:
        token = await oauth.hf.authorize_access_token(request)
        info = token.get('userinfo') or await oauth.hf.userinfo(token=token)
        identity = {'id': str(info['sub']), 'name': info.get('preferred_username') or info.get('name') or 'HF user'}
    except Exception:
        raise HTTPException(400, 'Sign-in did not complete. Please try again.')
    target = request.session.pop('return_poll', '')
    guest = request.session.get('voter', secrets.token_urlsafe(32))
    request.session.clear()
    request.session.update(user=identity, voter=guest)
    return RedirectResponse(BASE + '/#' + target, status_code=303)


@app.post('/api/logout')
def logout(request: Request):
    request.session.pop('user', None)
    return {'ok': True}


class NewPoll(BaseModel):
    question: str = Field(min_length=1, max_length=240)
    description: str = Field(default='', max_length=2000)
    options: list[str] = Field(min_length=2, max_length=12)


@app.post('/api/polls')
def create_poll(body: NewPoll, request: Request):
    identity = user(request)
    question = body.question.strip()
    options = [x.strip() for x in body.options]
    if not question or any(not x or len(x) > 160 for x in options) or len(set(x.casefold() for x in options)) != len(options):
        raise HTTPException(422, 'Add a question and 2–12 different options (up to 160 characters each).')
    with LOCK:
        if sum(p['owner_id'] == identity['id'] for p in state['polls'].values()) >= 100:
            raise HTTPException(422, 'You can have up to 100 polls. Delete an old poll first.')
        new = copy_state()
        poll_id = secrets.token_hex(8)
        new['polls'][poll_id] = dict(id=poll_id, question=question, description=body.description.strip(), options=options,
                                    owner_id=identity['id'], owner=identity['name'], created_at=now(), closed=False, votes=[])
        save(new)
    return {'id': poll_id, 'url': SHARE_BASE + '#' + poll_id}


@app.get('/api/polls')
def my_polls(request: Request):
    identity = user(request)
    with LOCK:
        return [dict(id=p['id'], question=p['question'], closed=p['closed'], total=len(p['votes']), created_at=p['created_at'])
                for p in reversed(list(state['polls'].values())) if p['owner_id'] == identity['id']]


@app.get('/api/polls/{poll_id}')
def poll_detail(poll_id: str, request: Request):
    with LOCK:
        p = get_poll(poll_id)
        keys = voter_keys(request)
        owner = request.session.get('user', {}).get('id') == p['owner_id']
        vote = next((v for v in p['votes'] if set(v['keys']) & set(keys)), None)
        return {**{k: p[k] for k in ['id', 'question', 'description', 'options', 'owner', 'created_at', 'closed']},
                'is_owner': owner, 'voted': vote is not None, 'choice': vote['choice'] if vote else None,
                'url': SHARE_BASE + '#' + p['id'],
                'results': [sum(v['choice'] == i for v in p['votes']) for i in range(len(p['options']))] if owner or vote is not None else None}


class Vote(BaseModel):
    choice: int


@app.post('/api/polls/{poll_id}/vote')
def vote(poll_id: str, body: Vote, request: Request):
    with LOCK:
        p = get_poll(poll_id)
        if p['closed']:
            raise HTTPException(409, 'Voting has closed.')
        keys = voter_keys(request)
        if any(set(v['keys']) & set(keys) for v in p['votes']):
            raise HTTPException(409, 'You have already voted in this poll.')
        if body.choice < 0 or body.choice >= len(p['options']):
            raise HTTPException(422, 'Choose one of the available options.')
        if len(p['votes']) >= 10000:
            raise HTTPException(409, 'This poll has reached its 10,000 vote limit.')
        new = copy_state()
        new['polls'][poll_id]['votes'].append({'choice': body.choice, 'keys': keys, 'at': now()})
        save(new)
    return {'ok': True}


class Status(BaseModel):
    closed: bool


@app.patch('/api/polls/{poll_id}')
def set_status(poll_id: str, body: Status, request: Request):
    with LOCK:
        owned(request, poll_id)
        new = copy_state()
        new['polls'][poll_id]['closed'] = body.closed
        save(new)
    return {'ok': True}


@app.delete('/api/polls/{poll_id}')
def delete_poll(poll_id: str, request: Request):
    with LOCK:
        owned(request, poll_id)
        new = copy_state()
        del new['polls'][poll_id]
        save(new)
    return {'ok': True}


@app.get('/api/polls/{poll_id}/export')
def export(poll_id: str, request: Request):
    with LOCK:
        p = owned(request, poll_id)
        stream = io.StringIO()
        writer = csv.writer(stream)
        writer.writerow(['Option', 'Votes'])
        for i, label in enumerate(p['options']):
            # Neutralize spreadsheet formulas in user-provided labels.
            writer.writerow(["'" + label if label[:1] in '=+-@\t\r' else label, sum(v['choice'] == i for v in p['votes'])])
    return Response(stream.getvalue(), media_type='text/csv', headers={'Content-Disposition': f'attachment; filename="poll-{poll_id}.csv"'})