tiny-snitch-prompt

#!/usr/bin/env python3
import logging
import base64
import time
import os
import sys
from PyQt5.QtGui import QColor
from PyQt5 import QtWidgets
from PyQt5 import QtCore
from PyQt5.QtCore import QThread, pyqtSignal
import argh
import subprocess

# slight delay before accepting user input to avoid accidental y/n response
delay_seconds = float(os.environ.get('TINYSNITCH_DELAY_SECONDS', '.5'))

durations = 'once', '1-minute', '3-minutes', '9-minutes', '24-hours', 'forever'
scopes = 'port-domain', 'domain'
actions = 'yes', 'no'

class MetaThread(QThread):
    signal = pyqtSignal('PyQt_PyObject')

    def __init__(self, src_port, dst_port):
        super().__init__()
        self.src_port = src_port
        self.dst_port = dst_port

    def run(self):
        meta = ''
        xs = subprocess.check_output(['ss', '-tupnH']).decode('utf-8').strip().splitlines()
        for x in xs:
            try:
                _proto, _state, _, _, _src, _dst, program = x.split()
            except ValueError:
                return
            else:
                if self.src_port == _src.split(':')[-1] and self.dst_port == _dst.split(':')[-1]:
                    pid = program.split('pid=')[-1].split(',')[0]
                    try:
                        path, *args = subprocess.check_output(f'ps --no-heading -o args {pid}', shell=True).decode('utf-8').split()
                    except subprocess.CalledProcessError:
                        path = program
                        args = ''
                    else:
                        args = ' '.join(args)
                    meta = f'[{pid}] {path} {args}'
                    break
        self.signal.emit(meta)

class MainWindow(QtWidgets.QMainWindow):
    def qt_setup(self):
        QtWidgets.QMainWindow.__init__(self)
        centralWidget = QtWidgets.QWidget(self)
        self.setCentralWidget(centralWidget)
        gridLayout = QtWidgets.QGridLayout(self)
        centralWidget.setLayout(gridLayout)
        palette = self.palette()
        palette.setColor(self.backgroundRole(), QColor(20, 20, 20))
        self.setPalette(palette)
        self.setStyleSheet("QLabel { font: %spt Mono; color: rgb(200,200,200); }" % os.environ.get('TINYSNITCH_FONTSIZE', '16'))
        self.title = QtWidgets.QLabel('', self)
        self.title.setAlignment(QtCore.Qt.AlignCenter)
        gridLayout.addWidget(self.title, 0, 0)

    def __init__(self, argv):
        self.qt_setup()
        self.start = time.monotonic()
        self.duration = '9-minute'
        self.subdomains = False
        self.port = True
        self.action = 'no'
        self.meta = None
        self.help = False
        self.legend = False
        self.reverse = False
        self.text = base64.b64decode(argv[0].encode()).decode().replace('<', '&lt;').replace('>', '&gt;').split(None, 1)
        self.seconds = 60
        self.timer = QtCore.QTimer()
        self.timer.setInterval(1000)
        self.timer.timeout.connect(self.tick)
        self.timer.start()
        self.update()

    def tick(self):
        self.seconds -= 1
        if self.seconds == 0:
            self.exit()
        self.update()

    def keyPressEvent(self, event, **kw):
        if time.monotonic() - self.start < delay_seconds:
            return
        key = event.key()
        if self.help or self.legend:
            self.help = self.legend = False
            self.update()
        else:
            if key == 89: # y
                self.action = 'yes'
                self.exit()
            elif key == 78: # n
                self.action = 'no'
                self.exit()
            elif key == 68: # d
                self.subdomains = not self.subdomains
            elif key == 80: # p
                self.port = not self.port
            elif key == 79: # o
                self.duration = 'once'
            elif key == 49: # 1
                self.duration = '1-minute'
            elif key == 50: # 2
                self.duration = '24-hour'
            elif key == 51: # 3
                self.duration = '3-minute'
            elif key == 57: # 9
                self.duration = '9-minute'
            elif key == 70: # f
                self.duration = 'forever'
            elif key == 72: # h
                self.help = True
            elif key == 76: # l
                self.legend = True
            elif key == 82: # r
                self.reverse = not self.reverse

        self.update()

    def finished(self, meta):
        self.meta = meta
        self.update()

    def update(self):
        if self.help:
            self.title.setText('''
                <p style="line-height: 50px;">
                    <span style="font-weight: bold;">actions</span><br/>(y)es<br/>(n)o<br/><br/>
                    <span style="font-weight: bold;">granularity</span><br/>(d)omain-toggle<br/>(p)ort-toggle<br/><br/>
                    <span style="font-weight: bold;">duration</span><br/>(o)nce<br/>(1)minute<br/>(3)minutes<br/>(9)minutes<br/>(2)4hours<br/>(f)orever<br/><br/>
                    <span style="font-weight: bold;">direction</span><br/>(r)reverse-src-dst<br/>
                </p>
            ''') # noqa
        elif self.legend:
            text = '\n\n'.join(['protocol', ' '.join(['src:port', '->', 'dst:port'])])
            text += '\n\n[pid] program args'
            self.title.setText(f'<span/>{text}\n\n{"-" * 250}\n\nduration subdomains\n\n(y)es (n)o (h)elp (l)egend\n\n{"-" * 250}\n\n{self.seconds} seconds until DENY'.replace('\n', '<br/>'))
        else:
            yellow = lambda x: '-' if x == '-' else f'<span style="color: rgb(220, 220, 0);">[{x}]</span>'
            proto, target = self.text
            src, arrow, dst = target.split()
            src_port = src.split(':')[-1]
            dst_port = dst.split(':')[-1]
            if self.reverse:
                src, dst = dst, src
                src_port, dst_port = dst_port, src_port
            if self.meta is None:
                self.thread = MetaThread(src_port, dst_port)
                self.thread.signal.connect(self.finished)
                self.thread.start()
            if proto.strip().endswith('-src'):
                a, b = src.split(':')
                src = f'{yellow(a)}:{b}'
                if self.port:
                    a, b = dst.split(':')
                    dst = f'{a}:{yellow(b)}'
            else:
                if self.subdomains and dst.split(':')[0].split('.')[-1].isalpha():
                    dst = '*.' + '.'.join(dst.split('.')[-2:])
                if not self.port:
                    a, b = dst.split(':')
                    dst = f'{yellow(a)}:{b}'
                else:
                    dst = yellow(dst)
            text = [proto, ' '.join([src, arrow, dst])]
            text = '\n\n'.join(text)
            if self.meta:
                text += f'\n\n{self.meta}'
            subdomains = "yes" if self.subdomains else "no"
            self.title.setText(f'<span/>{text}\n\n{"-" * 250}\n\n{self.duration} subdomains={subdomains}\n\n(y)es (n)o (h)elp (l)egend\n\n{"-" * 250}\n\n{self.seconds} seconds until DENY'.replace('\n', '<br/>'))

    def exit(self):
        self.update()
        self.action = {'yes': 'allow', 'no': 'deny'}[self.action]
        subdomains = "yes" if self.subdomains else "no"
        ports = "yes" if self.port else "no"
        reverse = "yes" if self.reverse else "no"
        print(f'{self.duration} {subdomains} {self.action} {ports} {reverse}')
        sys.exit(0)

@argh.dispatch_command
def main(*words):
    logging.basicConfig(level='INFO')
    app = QtWidgets.QApplication(sys.argv)
    win = MainWindow(words)
    win.show()
    app.exec_()