majuna/app/portal/origin.py

import urllib.parse
from datetime import datetime, timezone
from typing import List, Optional

import requests
import sqlalchemy
from flask import Blueprint, Response, flash, redirect, render_template, url_for
from flask.typing import ResponseReturnValue
from flask_wtf import FlaskForm
from sqlalchemy import exc
from wtforms import BooleanField, IntegerField, SelectField, StringField, SubmitField
from wtforms.validators import DataRequired

from app.extensions import db
from app.models.base import Group
from app.models.mirrors import Country, Origin
from app.portal.util import LifecycleForm, response_404, view_lifecycle

bp = Blueprint("origin", __name__)


class NewOriginForm(FlaskForm):  # type: ignore
    domain_name = StringField("Domain Name", validators=[DataRequired()])
    description = StringField("Description", validators=[DataRequired()])
    group = SelectField("Group", validators=[DataRequired()])
    auto_rotate = BooleanField("Enable auto-rotation?", default=True)
    smart_proxy = BooleanField("Requires smart proxy?", default=False)
    asset_domain = BooleanField("Used to host assets for other domains?", default=False)
    submit = SubmitField("Save Changes")


class EditOriginForm(FlaskForm):  # type: ignore[misc]
    description = StringField("Description", validators=[DataRequired()])
    group = SelectField("Group", validators=[DataRequired()])
    auto_rotate = BooleanField("Enable auto-rotation?")
    smart_proxy = BooleanField("Requires smart proxy?")
    asset_domain = BooleanField("Used to host assets for other domains?", default=False)
    risk_level_override = BooleanField("Force Risk Level Override?")
    risk_level_override_number = IntegerField(
        "Forced Risk Level", description="Number from 0 to 20", default=0
    )
    submit = SubmitField("Save Changes")


class CountrySelectForm(FlaskForm):  # type: ignore[misc]
    country = SelectField("Country", validators=[DataRequired()])
    submit = SubmitField("Save Changes", render_kw={"class": "btn btn-success"})


def final_domain_name(domain_name: str) -> str:
    session = requests.Session()
    r = session.get(f"https://{domain_name}/", allow_redirects=True, timeout=10)
    return urllib.parse.urlparse(r.url).netloc


@bp.route("/new", methods=["GET", "POST"])
@bp.route("/new/<group_id>", methods=["GET", "POST"])
def origin_new(group_id: Optional[int] = None) -> ResponseReturnValue:
    form = NewOriginForm()
    form.group.choices = [(x.id, x.group_name) for x in Group.query.all()]
    if form.validate_on_submit():
        origin = Origin()
        origin.group_id = form.group.data
        origin.domain_name = final_domain_name(form.domain_name.data)
        origin.description = form.description.data
        origin.auto_rotation = form.auto_rotate.data
        origin.smart = form.smart_proxy.data
        origin.assets = form.asset_domain.data
        origin.added = datetime.now(tz=timezone.utc)
        origin.updated = datetime.now(tz=timezone.utc)
        try:
            db.session.add(origin)
            db.session.commit()
            flash(f"Created new origin {origin.domain_name}.", "success")
            return redirect(url_for("portal.origin.origin_edit", origin_id=origin.id))
        except exc.SQLAlchemyError:
            flash("Failed to create new origin.", "danger")
            return redirect(url_for("portal.origin.origin_list"))
    if group_id:
        form.group.data = group_id
    return render_template("new.html.j2", section="origin", form=form)


@bp.route("/edit/<origin_id>", methods=["GET", "POST"])
def origin_edit(origin_id: int) -> ResponseReturnValue:
    origin: Optional[Origin] = Origin.query.filter(Origin.id == origin_id).first()
    if origin is None:
        return Response(
            render_template(
                "error.html.j2",
                section="origin",
                header="404 Origin Not Found",
                message="The requested origin could not be found.",
            ),
            status=404,
        )
    form = EditOriginForm(
        group=origin.group_id,
        description=origin.description,
        auto_rotate=origin.auto_rotation,
        smart_proxy=origin.smart,
        asset_domain=origin.assets,
        risk_level_override=origin.risk_level_override is not None,
        risk_level_override_number=origin.risk_level_override,
    )
    form.group.choices = [(x.id, x.group_name) for x in Group.query.all()]
    if form.validate_on_submit():
        origin.group_id = form.group.data
        origin.description = form.description.data
        origin.auto_rotation = form.auto_rotate.data
        origin.smart = form.smart_proxy.data
        origin.assets = form.asset_domain.data
        if form.risk_level_override.data:
            origin.risk_level_override = form.risk_level_override_number.data
        else:
            origin.risk_level_override = None
        origin.updated = datetime.now(tz=timezone.utc)
        try:
            db.session.commit()
            flash(f"Saved changes for origin {origin.domain_name}.", "success")
        except exc.SQLAlchemyError:
            flash("An error occurred saving the changes to the origin.", "danger")
    return render_template("origin.html.j2", section="origin", origin=origin, form=form)


@bp.route("/list")
def origin_list() -> ResponseReturnValue:
    origins: List[Origin] = Origin.query.order_by(Origin.domain_name).all()
    return render_template(
        "list.html.j2",
        section="origin",
        title="Web Origins",
        item="origin",
        new_link=url_for("portal.origin.origin_new"),
        items=origins,
        extra_buttons=[
            {
                "link": url_for("portal.origin.origin_onion"),
                "text": "Onion services",
                "style": "onion",
            }
        ],
    )


@bp.route("/onion")
def origin_onion() -> ResponseReturnValue:
    origins = Origin.query.order_by(Origin.domain_name).all()
    return render_template(
        "list.html.j2",
        section="origin",
        title="Onion Sites",
        item="onion service",
        new_link=url_for("portal.onion.onion_new"),
        items=origins,
    )


@bp.route("/destroy/<origin_id>", methods=["GET", "POST"])
def origin_destroy(origin_id: int) -> ResponseReturnValue:
    origin = Origin.query.filter(
        Origin.id == origin_id, Origin.destroyed.is_(None)
    ).first()
    if origin is None:
        return response_404("The requested origin could not be found.")
    return view_lifecycle(
        header=f"Destroy origin {origin.domain_name}",
        message=origin.description,
        success_message="All proxies from the destroyed origin will shortly be destroyed at their providers.",
        success_view="portal.origin.origin_list",
        section="origin",
        resource=origin,
        action="destroy",
    )


@bp.route("/country_remove/<origin_id>/<country_id>", methods=["GET", "POST"])
def origin_country_remove(origin_id: int, country_id: int) -> ResponseReturnValue:
    origin = Origin.query.filter(Origin.id == origin_id).first()
    if origin is None:
        return Response(
            render_template(
                "error.html.j2",
                section="origin",
                header="404 Pool Not Found",
                message="The requested origin could not be found.",
            ),
            status=404,
        )
    country = Country.query.filter(Country.id == country_id).first()
    if country is None:
        return Response(
            render_template(
                "error.html.j2",
                section="origin",
                header="404 Country Not Found",
                message="The requested country could not be found.",
            ),
            status=404,
        )
    if country not in origin.countries:
        return Response(
            render_template(
                "error.html.j2",
                section="origin",
                header="404 Country Not In Pool",
                message="The requested country could not be found in the specified origin.",
            ),
            status=404,
        )
    form = LifecycleForm()
    if form.validate_on_submit():
        origin.countries.remove(country)
        try:
            db.session.commit()
            flash("Saved changes to origin.", "success")
            return redirect(url_for("portal.origin.origin_edit", origin_id=origin.id))
        except sqlalchemy.exc.SQLAlchemyError:
            flash("An error occurred saving the changes to the origin.", "danger")
    return render_template(
        "lifecycle.html.j2",
        header=f"Remove {country.description} from the {origin.domain_name} origin?",
        message="Stop monitoring in this country.",
        section="origin",
        origin=origin,
        form=form,
    )


@bp.route("/country_add/<origin_id>", methods=["GET", "POST"])
def origin_country_add(origin_id: int) -> ResponseReturnValue:
    origin = Origin.query.filter(Origin.id == origin_id).first()
    if origin is None:
        return Response(
            render_template(
                "error.html.j2",
                section="origin",
                header="404 Origin Not Found",
                message="The requested origin could not be found.",
            ),
            status=404,
        )
    form = CountrySelectForm()
    form.country.choices = [
        (x.id, f"{x.country_code} - {x.description}") for x in Country.query.all()
    ]
    if form.validate_on_submit():
        country = Country.query.filter(Country.id == form.country.data).first()
        if country is None:
            return Response(
                render_template(
                    "error.html.j2",
                    section="origin",
                    header="404 Country Not Found",
                    message="The requested country could not be found.",
                ),
                status=404,
            )
        origin.countries.append(country)
        try:
            db.session.commit()
            flash("Saved changes to origin.", "success")
            return redirect(url_for("portal.origin.origin_edit", origin_id=origin.id))
        except sqlalchemy.exc.SQLAlchemyError:
            flash("An error occurred saving the changes to the origin.", "danger")
    return render_template(
        "lifecycle.html.j2",
        header=f"Add a country to {origin.domain_name}",
        message="Enable monitoring from this country:",
        section="origin",
        origin=origin,
        form=form,
    )