callfwdctl

#!/usr/bin/env python3

import socket
import array
import sys
import os
import tempfile
import datetime
import subprocess
import tarfile
import json
import argparse
import glob


CALLFWD_SERVER_PATH = "/run/callfwd"
if not os.path.exists(CALLFWD_SERVER_PATH):
    XDG_RUNTIME_DIR = os.getenv("XDG_RUNTIME_DIR")
    CALLFWD_SERVER_PATH = os.path.join(XDG_RUNTIME_DIR, "callfwd")


class CallFwdControl(object):
    def __init__(self, doorbell, server=None):
        if server is None:
            server = CALLFWD_SERVER_PATH
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        sock.bind(doorbell)
        sock.connect(server)
        self.sock = sock

    def _make_request(self, msg, fds):
        msg["stderr"] = len(fds)
        cdata = array.array("i", fds + [2])
        cmsg = (socket.SOL_SOCKET, socket.SCM_RIGHTS, cdata)
        self.sock.sendmsg([json.dumps(msg).encode('ascii')], [cmsg])

    def _wait_response(self):
        msg,_,_,_ = self.sock.recvmsg(10)
        print("Success" if msg == b'S' else 'Failure')
        if msg != b'S':
            exit(1)

    def _read_db_op(self, msg, path, row_size):
        msg["loaded"] = str(datetime.datetime.now())
        msg["file_name"] = path
        msg["stdin"] = 0
        if path.endswith(".tar.gz"):
            with tarfile.open(path) as tar:
                ti = tar.next()
                msg["inner_name"] = ti.name
                msg["row_estimate"] = ti.size // 23
            shell = ["tar", "xOf", path, msg["inner_name"]]
            with subprocess.Popen(shell, stdout=subprocess.PIPE) as p:
                self._make_request(msg, [p.stdout.fileno()])
        else:
            msg["row_estimate"] = os.stat(path).st_size // row_size
            with open(path, "r") as f:
                self._make_request(msg, [f.fileno()])
        self._wait_response()

    def reload_db(self, path, country, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "reload" }
        msg["country"] = country
        self._read_db_op(msg, path, 23)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def dnc_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "dnc_reload" }
        self._read_db_op(msg, path, 11)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def dno_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "dno_reload" }
        self._read_db_op(msg, path, 38)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def dno_npa_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "dno_npa_reload" }
        self._read_db_op(msg, path, 7)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def dno_npa_nxx_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "dno_npa_nxx_reload" }
        self._read_db_op(msg, path, 11)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def dno_npa_nxx_x_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "dno_npa_nxx_x_reload" }
        self._read_db_op(msg, path, 37)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def tollfree_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "tollfree_reload" }
        self._read_db_op(msg, path, 38)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def lerg_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "lerg_reload" }
        self._read_db_op(msg, path, 69)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def youmail_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "youmail_reload" }
        self._read_db_op(msg, path, 45)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def geo_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "geo_reload" }
        self._read_db_op(msg, path, 190)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def ftc_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "ftc_reload" }
        self._read_db_op(msg, path, 120)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def f404_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "404_reload" }
        self._read_db_op(msg, path, 51)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)
    
    def f606_reload_db(self, path, update):
        link_to = None
        if update is not None:
            candidates = sorted(glob.glob(update))
            if len(candidates) > 0:
                link_to = path
                path = candidates[-1]

        msg = { "cmd": "606_reload" }
        self._read_db_op(msg, path, 51)

        if link_to is not None:
            if os.path.exists(link_to):
                os.unlink(link_to)
            os.symlink(path, link_to)

    def verify_db(self, path, country):
        msg = { "cmd": "verify" }
        msg["country"] = country
        self._read_db_op(msg, path)

    def dump_db(self, path, country):
        msg = { "cmd": "dump" }
        msg["file_name"] = path
        msg["country"] = country
        msg["stdout"] = 0
        with open(path, "w+") as f:
            self._make_request(msg, [f.fileno()])
        self._wait_response()

    def reload_acl(self, acl):
        msg = { "cmd": "acl" }
        msg["file_name"] = acl
        msg["stdin"] = 0
        with open(acl, "r") as f:
            self._make_request(msg, [f.fileno()])
        self._wait_response()

    def status(self):
        msg = { "cmd": "meta" }
        self._make_request(msg, [])
        self._wait_response()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-S', '--socket', type=str, default=CALLFWD_SERVER_PATH,
                        help="Path to server control socket")
    subparsers = parser.add_subparsers()

    reload_group = subparsers.add_parser('reload')
    reload_group.add_argument('-c', '--country', type=str, default='US',
                              help="Country code (US, CA)")
    reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    reload_group.add_argument('db', type=str, help="Path to database")
    reload_group.set_defaults(func=CallFwdControl.reload_db)
    reload_group.set_defaults(args=['db', 'country', 'update'])

    dnc_reload_group = subparsers.add_parser('dnc_reload')
    dnc_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    dnc_reload_group.add_argument('db', type=str, help="Path to database")
    dnc_reload_group.set_defaults(func=CallFwdControl.dnc_reload_db)
    dnc_reload_group.set_defaults(args=['db', 'update'])

    dno_reload_group = subparsers.add_parser('dno_reload')
    dno_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    dno_reload_group.add_argument('db', type=str, help="Path to database")
    dno_reload_group.set_defaults(func=CallFwdControl.dno_reload_db)
    dno_reload_group.set_defaults(args=['db', 'update'])

    dno_npa_reload_group = subparsers.add_parser('dno_npa_reload')
    dno_npa_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    dno_npa_reload_group.add_argument('db', type=str, help="Path to database")
    dno_npa_reload_group.set_defaults(func=CallFwdControl.dno_npa_reload_db)
    dno_npa_reload_group.set_defaults(args=['db', 'update'])

    dno_npa_nxx_reload_group = subparsers.add_parser('dno_npa_nxx_reload')
    dno_npa_nxx_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    dno_npa_nxx_reload_group.add_argument('db', type=str, help="Path to database")
    dno_npa_nxx_reload_group.set_defaults(func=CallFwdControl.dno_npa_nxx_reload_db)
    dno_npa_nxx_reload_group.set_defaults(args=['db', 'update'])

    dno_npa_nxx_x_reload_group = subparsers.add_parser('dno_npa_nxx_x_reload')
    dno_npa_nxx_x_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    dno_npa_nxx_x_reload_group.add_argument('db', type=str, help="Path to database")
    dno_npa_nxx_x_reload_group.set_defaults(func=CallFwdControl.dno_npa_nxx_x_reload_db)
    dno_npa_nxx_x_reload_group.set_defaults(args=['db', 'update'])

    tollfree_reload_group = subparsers.add_parser('tollfree_reload')
    tollfree_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    tollfree_reload_group.add_argument('db', type=str, help="Path to database")
    tollfree_reload_group.set_defaults(func=CallFwdControl.tollfree_reload_db)
    tollfree_reload_group.set_defaults(args=['db', 'update'])

    lerg_reload_group = subparsers.add_parser('lerg_reload')
    lerg_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    lerg_reload_group.add_argument('db', type=str, help="Path to database")
    lerg_reload_group.set_defaults(func=CallFwdControl.lerg_reload_db)
    lerg_reload_group.set_defaults(args=['db', 'update'])

    youmail_reload_group = subparsers.add_parser('youmail_reload')
    youmail_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    youmail_reload_group.add_argument('db', type=str, help="Path to database")
    youmail_reload_group.set_defaults(func=CallFwdControl.youmail_reload_db)
    youmail_reload_group.set_defaults(args=['db', 'update'])
    
    geo_reload_group = subparsers.add_parser('geo_reload')
    geo_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    geo_reload_group.add_argument('db', type=str, help="Path to database")
    geo_reload_group.set_defaults(func=CallFwdControl.geo_reload_db)
    geo_reload_group.set_defaults(args=['db', 'update'])

    ftc_reload_group = subparsers.add_parser('ftc_reload')
    ftc_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    ftc_reload_group.add_argument('db', type=str, help="Path to database")
    ftc_reload_group.set_defaults(func=CallFwdControl.ftc_reload_db)
    ftc_reload_group.set_defaults(args=['db', 'update'])

    f404_reload_group = subparsers.add_parser('404_reload')
    f404_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    f404_reload_group.add_argument('db', type=str, help="Path to database")
    f404_reload_group.set_defaults(func=CallFwdControl.f404_reload_db)
    f404_reload_group.set_defaults(args=['db', 'update'])

    f606_reload_group = subparsers.add_parser('606_reload')
    f606_reload_group.add_argument('-u', '--update', type=str, default=None,
                              help="A directory where to search for updates")
    f606_reload_group.add_argument('db', type=str, help="Path to database")
    f606_reload_group.set_defaults(func=CallFwdControl.f606_reload_db)
    f606_reload_group.set_defaults(args=['db', 'update'])

    verify_group = subparsers.add_parser('verify')
    verify_group.add_argument('-c', '--country', type=str, default='US',
                              help="Country code (US, CA)")
    verify_group.add_argument('db', type=str, help="Path to database")
    verify_group.set_defaults(func=CallFwdControl.verify_db)
    verify_group.set_defaults(args=['db', 'country'])

    dump_group = subparsers.add_parser('dump')
    dump_group.add_argument('-c', '--country', type=str, default='US',
                            help="Country code (US, CA)")
    dump_group.add_argument('db', type=str, help="Path to database")
    dump_group.set_defaults(func=CallFwdControl.dump_db)
    dump_group.set_defaults(args=['db', 'country'])

    acl_group = subparsers.add_parser('acl')
    acl_group.add_argument('csv', type=str, help="Path to PGSQL dump")
    acl_group.set_defaults(func=CallFwdControl.reload_acl)
    acl_group.set_defaults(args=['csv'])

    status_group = subparsers.add_parser('status')
    status_group.set_defaults(func=CallFwdControl.status)
    status_group.set_defaults(args=[])

    options = parser.parse_args()

    with tempfile.TemporaryDirectory('callfwdctl') as tmpdir:
        doorbell = os.path.join(tmpdir, 'doorbell')
        ctl = CallFwdControl(doorbell, options.socket)
        options.func(ctl, *map(options.__getattribute__, options.args))