#!/usr/bin/env python3
# This file is part of OpenSnitch.
#
# Copyright(c) 2017 Simone Margaritelli
# evilsocket@gmail.com
# http://www.evilsocket.net
#
# This file may be licensed under the terms of of the
# GNU General Public License Version 2 (the ``GPL'').
#
# Software distributed under the License is distributed
# on an ``AS IS'' basis, WITHOUT WARRANTY OF ANY KIND, either
# express or implied. See the GPL for the specific language
# governing rights and limitations.
#
# You should have received a copy of the GPL along with this
# program. If not, go to http://www.gnu.org/licenses/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from argparse import ArgumentParser
from os.path import expanduser
import warnings
import logging
import prctl
import pwd
import sys
import os

# Ensure Qt4 wont be loaded by matplotlib
try:
    import matplotlib as mpl
except ImportError:
    pass
else:
    mpl.rcParams['backend'] = 'Qt5Agg'

from opensnitch.version import VERSION
from opensnitch.snitch import Snitch


REQUIRED_CAPS = ((prctl.CAP_NET_RAW, prctl.ALL_FLAGS, True),
                 (prctl.CAP_DAC_OVERRIDE, prctl.ALL_FLAGS, True),
                 (prctl.CAP_NET_ADMIN, prctl.ALL_FLAGS, True))


parser = ArgumentParser()
parser.add_argument("--log-file", dest="logfile", default=None,
                    help="Log to file", metavar="FILE")
parser.add_argument("--debug", dest="debug",
                    action="store_true", default=False,
                    help="Enable debug logs")
parser.add_argument("--database",  dest="database", default=None,
                    help="Database path", metavar="FILE")
parser.add_argument("--user",  dest="user",
                    default=os.environ.get('SUDO_USER', 'root'),
                    help="Run as user")


# At some point Scapy devs will realize how bothering their fucking warnings
# are while importing scapy.all ...
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
warnings.filterwarnings("ignore", category=RuntimeWarning, module="gtk")


if __name__ == '__main__':
    args = parser.parse_args()

    logging.basicConfig(
        format='[%(asctime)s] (%(levelname)s) %(message)s',
        level=logging.INFO if args.debug is False else logging.DEBUG,
        filename=args.logfile)

    if not os.geteuid() == 0:
        sys.exit('OpenSnitch must be run as root.')

    if 'DBUS_SESSION_BUS_ADDRESS' not in os.environ:
        raise RuntimeError('DBUS_SESSION_BUS_ADDRESS not set')

    # set_keepcaps allows us to keep caps across setuid call
    prctl.set_keepcaps(True)
    prctl.set_caps(*REQUIRED_CAPS)

    user = pwd.getpwnam(args.user)
    os.setgid(user.pw_gid)
    os.setuid(user.pw_uid)
    # The effective capability set is cleared on setuid
    # Re-set the effective set
    prctl.set_caps(*REQUIRED_CAPS)

    dbfile = args.database or os.path.join(
        expanduser("~%s" % args.user), "opensnitch.db")

    snitch = Snitch(dbfile)
    try:
        logging.info("OpenSnitch v%s running with pid %d.",
                     VERSION, os.getpid())
        snitch.start()
    except KeyboardInterrupt as e:
        pass
    finally:
        logging.info("Quitting ...")
        snitch.stop()

        # Temporary hack to handle Ctrl-C
        os.kill(os.getpid(), 15)
