#!/usr/bin/python3
import apt
import os
import csv
#polkit check for debian 
def check_debian(dist):

    if "stretch" in dist[1]:
        return "0.105-18+deb9u2"
    elif "buster" in dist[1]:
        return "0.105-25+deb10u1"
    elif "bullseye" in dist[1]:
        return "0.105-31+deb11u1"

# Return True is vulnerable, False if not
def check_deb_varients(dist):
    
    cache = apt.cache.Cache()
    cache.open()

    pkg = cache["policykit-1"]
    pkg_ver = pkg.installed.version

    fixed_ver = None

    if "Debian" in dist[0]:
        fixed_ver = check_debian(dist)

    return pkg_ver < fixed_ver
def fix():
    os.system("sudo dpkg -i /mnt/fix_polkit_CVE-2021-4034.deb")
    print("[-] Please reboot and re-run the patch to verify")
    os.system("sudo touch /var/log/polkit.txt && echo 'patch file executed $(date)'")
def find_distro():

    RELEASE_DATA = {}
    
    with open("/etc/os-release") as f:
        reader = csv.reader(f, delimiter="=")
        for row in reader:
            if row:
                RELEASE_DATA[row[0]] = row[1]

    if RELEASE_DATA["ID"] in ["debian", "raspbian"]:
        with open("/etc/debian_version") as f:
            DEBIAN_VERSION = f.readline().strip()

        major_version = DEBIAN_VERSION.split(".")[0]
        version_split = RELEASE_DATA["VERSION"].split(" ", maxsplit=1)

        if version_split[0] == major_version:
            # Just major version shown, replace it with the full version
            RELEASE_DATA["VERSION"] = " ".join([DEBIAN_VERSION] + version_split[1:])

    return (RELEASE_DATA["NAME"], RELEASE_DATA["VERSION"])

def main():
    
    banner = """---> polkit vurln scanner <---

This test is currently working on Debian (stretch, buster, and bullseye) only
If you face code errors reachout @ ayush.bhatia@coraltele.com"""
    
    print(banner)
    print()

    dist = find_distro()

    if "Debian" in dist[0]:
        print("[*] Test started")
        is_vuln = False
        is_vuln = check_deb_varients(dist)
    
        if is_vuln:
            print("[-] Your polkit package is vulnerable.\nUpdate it using: apt install policykit-1")
            print("[-] We will be auto updating the package to fix the issue")
            fix()
        else:
            print("[+] Your polkit package is not vulnerable. Keep being secure")
    
    else:
        print("[-] This test is currently working on Debian only.")



if __name__ == "__main__":

    main()