veil-server/badsites.py

#!/usr/bin/python3
# Author: Farooq Karimi Zadeh
# Email: fkz@riseup.net
# See LICENSE for licence information
import pickle
import cherrypy # <3
import random
import os
from hashlib import sha512 # for password hashing
import re
import lmdb # <3
from bsconfig import BSConfig
SITE_RE = re.compile("([\w][^\/=\s]+)\/?|(^w{3}[\.\w][^\/\=\s]{2,})")
# ^ I use this RE to determine if a valid url has been sent
ENV = None # LMDB Environment


@cherrypy.tools.register("before_handler")
def auth(groups):
    global ENV
    sess = cherrypy.session
    if "login?" in sess:
        if sess["login?"] in groups:
            return
        else:
            raise cherrypy.HTTPRedirect("/permission")
    else:
        raise cherrypy.HTTPRedirect("/login")
    
class Veil:
    def __init__(self, config):
        self.config = config
        self.templates = dict()
        self.env = ENV
        self.img_db = self.env.open_db(b"imgs")
        self.page_db = self.env.open_db(b"pages")
        self.admin_db = self.env.open_db(b"admins")
        self.img_exts = ("apng", "bmp", "gif",
                        "ico", "cur", "jpg",
                        "jpeg", "jfif", "pjpeg",
                        "pjp", "png", "svg",
                        "tif", "tiff", "webp")
                    
        for template in os.listdir(config.templates_path):
            with open(config.templates_path + "/" + template) as fp:
                self.templates[template.replace(".html", "")] = fp.read()

    def is_img(self, addr):
        for ext in self.img_exts:
            if addr.endswith("." + ext):
                return True
        return False
                    
    def __vote1(self, addr, ip):
        if self.is_img(addr):
            db = self.img_db
        else:
            db = self.page_db

        with self.env.begin(db=db, write=True) as txn:
            data = txn.get(addr)
            if data:
                data = pickle.loads(data)
            else:
                data = set()
            if type(data) != set:
                return True
            data.add(ip)
            txn.replace(addr.encode(), pickle.dumps(data))

    def __unvote1(self, addr, ip):
        if self.is_img(addr):
            db = self.img_db
        else:
            db = self.page_db

        with self.env.begin(write=True, db=db) as txn:
            data = txn.get(addr)
            if data:
                data = pickle.loads(data)
                if type(data) != set:
                    return
                data.discard(ip)
                txn.replace(addr.encode(), data)

    @cherrypy.expose
    def permission(self):
        return "Sorry! Permission Denied!"
    
    @cherrypy.expose
    @cherrypy.tools.caching(delay=1)
    @cherrypy.tools.json_out()
    def index(self):
        result = dict()
        result["pages"] = dict()
        result["imgs"] = dict()
        
        with self.env.begin(db = self.page_db) as txn:
            for addr, ips in txn.cursor():
                result["pages"][addr] = len(ips) if type(ips) == set else -1
        with self.env.begin(db = self.img_db) as txn:
            for addr, ips in txn.cursor():
                result["imgs"][addr] = len(ips) if type(ips) == set else -1

        return result

    @cherrypy.expose
    def vote(self, addr):
        if addr.startswith("https"):
            addr = addr.replace("https", "http", 1)

        if addr.startswith("http"):
            self.__vote1(addr, cherrypy.request.remote.ip)
            return "Vote for the page successfully recorded"
        else:
            self.__vote1(addr, cherrypy.request.remote.ip)
            return "Vote for the page successfully recorded"

    @cherrypy.expose
    def unvote(self, addr):
        if addr.startswith("https"):
            addr = addr.replace("https", "http", 1)

        if addr.startswith("http"):
            self.__unvote1(addr, cherrypy.request.remote.ip)
            return "Vote for the page successfully removed"
        else:
            self.__unvote1(addr, cherrypy.request.remote.ip)
            return "Vote for the site successfully removed"

    @cherrypy.expose
    @cherrypy.tools.auth(groups=["tomato", "admin"])
    def addnew(self, data=None):
        if not data:
            return "Empty field: data"
        
        data = data.replace("\r", "")

        for line in data.split("\n"):
            line = line.strip()
            if not line:
                continue
            if line.startswith("https"):
                line = line.replace("https", "http", 1)

            if self.is_img(line):
                db = self.img_db
            else:
                db = self.page_db

            with self.env.begin(write=True, db=db) as txn:
                if txn.get(line.encode()):
                    txn.replace(line.encode(), pickle.dumps(None))
                else:
                    txn.put(line.encode(), pickle.dumps(None))

        return "Done :)"
            
    @cherrypy.expose
    @cherrypy.tools.auth(groups=["tomato", "admin"])
    def delete(self, addr):
        addr = addr.encode()
        if self.is_img(addr):
            db = self.img_db
        else:
            db = self.page_db

        with self.env.begin(write=True, db=db) as txn:
            if txn.delete(addr):
                return "Deleted!"
            else:
                return "No such thing in DB"
    
    @cherrypy.expose
    @cherrypy.tools.auth(groups=["tomato", "admin"])
    def panel(self):
        return self.templates["panel"]

    @cherrypy.expose
    def logout(self):
        if "login?" in cherrypy.session:
            cherrypy.session["login?"] = None
            raise cherrypy.HTTPRedirect("/login")
        else:
            return "You have not logged in..."

    @cherrypy.expose
    @cherrypy.tools.auth(groups=["tomato"])
    def addadmin(self, username="", password=""):
        if not (username and password):
            return self.templates["addadmin"]
        username = username.encode()
        password = password.encode()
        
        with self.env.begin(write=True, db=self.admin_db) as txn:
            if txn.get(username):
                txn.replace(username, password)
            else:
                txn.put(username, password)

        return "Well done :)"
    
    @cherrypy.expose
    def login(self, username="", password=""):
        if not(username and password):
            return self.templates["login"]
        if len(password) > 1024 or len(username) > 1024:
            return "Too long username or password"
        
        pass_hash = sha512(password.encode()).hexdigest()
        if username == "tomato" and pass_hash == self.config.tomato_pwd_hash:
            cherrypy.session["login?"] = "tomato"
            return "Welcome Tomato!"
        with self.env.begin(db=self.admin_db) as txn:
            the_hash = txn.get(username)
            if the_hash and the_hash.decode() == pass_hash:
                sess = cherrypy.session
                sess["login?"] = username
                raise cherrypy.HTTPRedirect("/panel")
            else:
                return "Incorrect username of password."

if __name__ == "__main__":
    conf = {
        "/": {
            "tools.staticdir.root": os.path.abspath(os.getcwd()),
            "tools.sessions.on": True
        },
    }
    cherrypy.config.update({
                    "server.socket_host": BSConfig.bindto,
                    "server.socket_port": BSConfig.port})
    ENV = lmdb.open(BSConfig.db_path, max_dbs=3)
    cherrypy.quickstart(Veil(BSConfig), "/", conf)