#! /usr/bin/python3
# SPDX-License-Identifier: MIT

import os
import sys
import re
import subprocess
import time
import ipaddress
import argparse
import tempfile
import tabulate
import dbus

bus = None
manager = None

def fatal(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)
    sys.exit(1)

def set_property_wait(interface, prop, value, timeout=10):
    interface.SetProperty(prop, value, timeout=120)
    for _i in range(0, timeout):
        state = interface.GetProperties()
        if state[prop] == value:
            return True
        time.sleep(1)

    return False

def init():
    global bus, manager
    try:
        bus = dbus.SystemBus()
    except Exception as e:
        fatal(e)

    try:
        manager = dbus.Interface(bus.get_object('org.ofono', '/'), 'org.ofono.Manager')
    except dbus.exceptions.DBusException:
        fatal("Could not aquire org.ofono.Manager on dbus")

def action_list():
    init()
    global manager, bus
    modems = manager.GetModems()

    if len(modems) == 0:
        print("No modems found")
        return

    modem_info = []
    context_info = []

    for path, properties in modems:
        model = path[1:]
        powered = properties.get('Powered', False)
        online = properties.get('Online', False)

        if not powered:
            modem_info.append([model, "Unpowered", "N/A", "N/A"])
            continue

        if not online:
            modem_info.append([model, "Offline", "N/A", "N/A"])
            continue

        try:
            registration_interface = dbus.Interface(bus.get_object('org.ofono', path), 'org.ofono.NetworkRegistration')
            properties = registration_interface.GetProperties()
            status = str(properties["Status"])
            network_name = str(properties["Name"])
            if status == "registered":
                if "Strength" in properties:
                    strength = float(properties["Strength"])
                else:
                    strength = 0.0
                registration = f"Registered to {network_name} ({strength}%)"
            else:
                registration = f"{status.capitalize()}"
            technology = str(properties.get("Technology", "Unknown"))
        except dbus.exceptions.DBusException:
            registration = "Unregistered"
            network_name = "Unknown"
            technology = "Unknown"

        try:
            sim_manager = dbus.Interface(bus.get_object('org.ofono', path), 'org.ofono.SimManager')
            properties = sim_manager.GetProperties()
            if properties['Present']:
                sim = properties.get('ServiceProviderName', '')
                if not sim:
                    sim = network_name
            else:
                sim = "No SIM"
        except dbus.exceptions.DBusException:
            sim = network_name

        try:
            ims = dbus.Interface(bus.get_object('org.ofono', path), 'org.ofono.IpMultimediaSystem')
            ims_properties = ims.GetProperties()
            ims_status = "Registered" if ims_properties.get('Registered', False) else "Not Registered"
            voice_capable = "V" if ims_properties.get('VoiceCapable', False) else "-"
            sms_capable = "S" if ims_properties.get('SmsCapable', False) else "-"
            ims_info = f"{ims_status} ({voice_capable}{sms_capable})"
        except dbus.exceptions.DBusException:
            ims_info = "Unknown"

        modem_info.append([model, registration, sim, ims_info, technology])

        try:
            connection_manager = dbus.Interface(bus.get_object('org.ofono', path), 'org.ofono.ConnectionManager')
            contexts = connection_manager.GetContexts()
            for context in contexts:
                context_path = context[0]
                apn = context[1].get('AccessPointName', 'N/A')
                settings = context[1].get('Settings', {})
                interface = settings.get('Interface', 'N/A')
                method = settings.get('Method', 'N/A')
                address = settings.get('Address', 'N/A')
                gateway = settings.get('Gateway', 'N/A')
                dns = ', '.join(settings.get('DomainNameServers', []))
                context_info.append([model, context_path, apn, interface, method, address, gateway, dns])
        except dbus.exceptions.DBusException:
            context_info.append([model, "N/A", "N/A", "N/A", "N/A", "N/A", "N/A", "N/A"])

    modem_table = tabulate.tabulate(
        modem_info,
        headers=["Modem", "Status", "SIM", "IMS Status (VS)", "Access Technology"],
        tablefmt="grid"
    )

    context_table = tabulate.tabulate(
        context_info,
        headers=["Modem", "Context Path", "APN", "Interface", "Method", "Address", "Gateway", "DNS"],
        tablefmt="grid"
    )

    print("Modem Information:")
    print(modem_table)
    print("\nContext Information:")
    print(context_table)

def action_power(component, state, command):
    message = {
        'poweron': ["Powered on {}", "Could not power on {}"],
        'poweroff': ["Powered off {}", "Could not power off {}"],
        'online': ["Brought {} online", "Could not online {}"],
        'offline': ["Took {} offline", "Could not offline {}"]
    }

    init()
    global manager, bus
    modems = manager.GetModems()

    if len(modems) == 0:
        print("No modems found")
        sys.exit(1)

    if component == 'Online' and state:
        powered = modems[0][1]['Powered'] == 1
        if not powered:
            print("Trying to online a modem that's not powered on. Running power on first...")
            action_power('Powered', True, 'poweron')

    for path in modems:
        model = path[0]
        modem = dbus.Interface(bus.get_object('org.ofono', model), 'org.ofono.Modem')
        if set_property_wait(modem, component, dbus.Boolean(1 if state else 0)):
            print(message[command][0].format(model))
            return
        fatal(message[command][1].format(model))

def action_scan_operators():
    init()
    global manager, bus
    modems = manager.GetModems()
    if len(modems) == 0:
        print("No modems found")
        sys.exit(1)
    modem = modems[0][0]

    netreg = dbus.Interface(bus.get_object('org.ofono', modem), 'org.ofono.NetworkRegistration')

    print("Scanning for operators... (100 seconds)")

    operators = netreg.Scan(timeout=100)
    result = []
    for _, properties in operators:
        tech = ", ".join(list(properties['Technologies']))
        result.append([properties['Name'], properties['Status'], tech, properties['MobileCountryCode']])

    print(tabulate.tabulate(result, headers=['Operator', 'Status', 'Technology', 'MCC']))

def action_wan(connect=False, disconnect=False, resolv=False, context_number=None):
    init()
    global manager, bus
    modems = manager.GetModems()
    if len(modems) == 0:
        print("No modems found")
        sys.exit(1)
    modem = modems[0][0]

    if 'Powered' in modems[0][1] and not modems[0][1]['Powered']:
        print("The modem is not powered, can't control WAN settings")
        print("You can power on the modem using ofonoctl poweron")
        sys.exit(1)

    if 'Online' in modems[0][1] and not modems[0][1]['Online']:
        print("The modem is offline, can't control WAN settings")
        print("You can bring the modem online using ofonoctl online")
        sys.exit(1)

    connman = dbus.Interface(bus.get_object('org.ofono', modem), 'org.ofono.ConnectionManager')
    try:
        contexts = connman.GetContexts()
    except dbus.exceptions.DBusException:
        print("Could not fetch contexts on the modem")
        sys.exit(1)

    if (connect or disconnect) and context_number is None:
        context_number = 1  # Default to context 1 for connect/disconnect if not specified

    if context_number is not None:
        if context_number < 1 or context_number > len(contexts):
            print(f"Invalid context number. Available contexts: 1 to {len(contexts)}")
            sys.exit(1)
        context_to_modify = contexts[context_number - 1]
    else:
        context_to_modify = None

    result = []
    has_flushed = False
    dns_servers = []

    def process_settings(settings, ip_version):
        nonlocal has_flushed, result, dns_servers
        if isinstance(settings, dbus.Dictionary) and ("Method" in settings or "Address" in settings):
            s = dict(settings)
            interface = s.get('Interface')

            if connect and not has_flushed and interface:
                cmd = ['ip', 'addr', 'flush', 'dev', interface]
                try:
                    subprocess.check_output(cmd)
                    has_flushed = True
                except subprocess.CalledProcessError as e:
                    print(f"Failed to flush: {e}")
                    return

            if ip_version == "ipv4":
                method = s.get("Method", "")
                address = s.get("Address", "") if method == "static" else ""
                gateway = s.get("Gateway", "") if method == "static" else ""
                dns = ", ".join(s.get("DomainNameServers", [])) if method == "static" else ""
                if address:
                    address += "/" + str(ipaddress.IPv4Network(f'0.0.0.0/{s["Netmask"]}').prefixlen)
                dns_servers += s.get("DomainNameServers", [])
            else:
                method = ""
                address = s.get("Address", "")
                gateway = s.get("Gateway", "")
                dns = ", ".join(s.get("DomainNameServers", []))
                dns_servers += s.get("DomainNameServers", [])

            result.append([context_number, interface, ip_version, properties.get("AccessPointName", ""), method, address, gateway, dns])

            if connect and interface:
                if address:
                    cmd = ['ip', 'addr', 'add', address, 'dev', interface]
                    try:
                        subprocess.check_output(cmd)
                    except subprocess.CalledProcessError as e:
                        print(f"Failed to set the ip address {address}: {e}")
                        return

                if gateway:
                    cmd = ['ip', 'route', 'replace', 'default', 'via', gateway, 'dev', interface]
                    try:
                        subprocess.check_output(cmd)
                    except subprocess.CalledProcessError as e:
                        print(f"Failed to set the default gateway {gateway}: {e}")

                for dns_server in s.get("DomainNameServers", []):
                    cmd = ['ip', 'route', 'replace', f"{dns_server}/{'32' if ip_version == 'ipv4' else '128'}", 'via', gateway, 'dev', interface]
                    try:
                        subprocess.check_output(cmd)
                    except subprocess.CalledProcessError as e:
                        print(f"Failed to set the default dns {dns_server} for gateway {gateway}: {e}")

    if connect or disconnect:
        if context_to_modify:
            context_path, properties = context_to_modify
            print(f"{'Connecting' if connect else 'Disconnecting'} context {context_number}")

            try:
                new_state = dbus.Boolean(1) if connect else dbus.Boolean(0)
                context_interface = dbus.Interface(bus.get_object('org.ofono', context_path),
                                                   'org.ofono.ConnectionContext')
                context_interface.SetProperty("Active", new_state)
                print(f"{'Connected' if connect else 'Disconnected'} context {context_number}")
            except dbus.exceptions.DBusException as e:
                print(f"Failed to {'connect' if connect else 'disconnect'} context {context_number}: {e}")
                return

            if connect:
                time.sleep(2) # wait a bit for context to activate and get an ip

                try:
                    updated_contexts = connman.GetContexts()
                except dbus.exceptions.DBusException:
                    print("Could not fetch updated contexts after activation")
                    return

                activated_context = next((ctx for ctx in updated_contexts if ctx[0] == context_path), None)

                if activated_context:
                    _, updated_properties = activated_context
                    settings4 = updated_properties.get('Settings', {})
                    settings6 = updated_properties.get('IPv6.Settings', {})

                    process_settings(settings4, "ipv4")
                    process_settings(settings6, "ipv6")

                    if resolv:
                        update_resolvconf(dns_servers)

                    if result:
                        headers = ["Context", "Interface", "Protocol", "APN", "Method", "Address", "Gateway", "DNS"]
                        print("Updated context information:")
                        print(tabulate.tabulate(result, headers=headers, tablefmt="simple"))
                else:
                    print("Could not find the activated context in the updated contexts")
            else:
                print(f"Context {context_number} disconnected")

    else:
        all_contexts = []
        for i, (_, props) in enumerate(contexts, 1):
            settings = props.get('Settings', {})
            ipv4 = dict(settings)
            all_contexts.append([
                i,
                ipv4.get('Interface', ''),
                'ipv4',
                props.get('AccessPointName', ''),
                ipv4.get('Method', ''),
                f"{ipv4.get('Address', '')}/{ipv4.get('Netmask', '')}" if ipv4.get('Address') else '',
                ipv4.get('Gateway', ''),
                ', '.join(ipv4.get('DomainNameServers', []))
            ])

            ipv6 = dict(props.get('IPv6.Settings', {}))
            if ipv6:
                all_contexts.append([
                    i,
                    ipv6.get('Interface', ''),
                    'ipv6',
                    props.get('AccessPointName', ''),
                    '',
                    ipv6.get('Address', ''),
                    ipv6.get('Gateway', ''),
                    ', '.join(ipv6.get('DomainNameServers', []))
                ])
        if all_contexts:
            headers = ["Context", "Interface", "Protocol", "APN", "Method", "Address", "Gateway", "DNS"]
            print(tabulate.tabulate(all_contexts, headers=headers, tablefmt="simple"))
        else:
            print("No contexts found")

def action_sms(destination, message=None):
    init()
    global manager, bus
    modems = manager.GetModems()
    if len(modems) == 0:
        print("No modems found")
        sys.exit(1)

    modem = modems[0][0]
    if message is None:
        editor = 'nano'
        if 'EDITOR' in os.environ:
            editor = os.environ['EDITOR']
        if 'VISUAL' in os.environ:
            editor = os.environ['VISUAL']
        with tempfile.NamedTemporaryFile(suffix='.txt', prefix='sms-') as buffer:
            subprocess.call([editor, buffer.name])
            buffer.seek(0)
            message = buffer.read().decode().strip()

    if len(message) == 0:
        print("Message empty. Aborting...")
        sys.exit(1)

    mm = dbus.Interface(bus.get_object('org.ofono', modem), 'org.ofono.MessageManager')
    mm.SendMessage(destination, message)
    print("Sent")

def action_sms_get():
    init()
    global manager, bus
    modems = manager.GetModems()
    if len(modems) == 0:
        print("No modems found")
        sys.exit(1)
    modem = modems[0][0]

    mm = dbus.Interface(bus.get_object('org.ofono', modem), 'org.ofono.MessageManager')
    messages = mm.GetMessages()
    print(messages)

def update_resolvconf(nameservers):
    with open('/etc/resolv.conf', encoding='utf-8') as handle:
        current = handle.read()

    header = 'DNS servers set by ofonoctl'

    regex = rf"# {header}.+# end\n"
    new_block = f'# {header}\n'
    for ns in nameservers:
        new_block += f'nameserver {ns}\n'
    new_block += '# end\n'

    if header in current:
        new_file = re.sub(regex, new_block, current, flags=re.MULTILINE | re.DOTALL)
    else:
        new_file = current + '\n' + new_block

    with open('/etc/resolv.conf', 'w', encoding='utf-8') as handle:
        handle.write(new_file)

def main():
    parser = argparse.ArgumentParser(description="Ofono control tool")
    sub = parser.add_subparsers(title="action", dest="action")
    sub.add_parser('list', help="List modems")
    sub.add_parser('poweron', help="Enable power to modem")
    sub.add_parser('poweroff', help="Disable power to modem")
    sub.add_parser('online', help="Enable modem")
    sub.add_parser('offline', help="Disable modem")
    sub.add_parser('operators', help="Display operator info")

    parser_wan = sub.add_parser('wan', help="Control internet access")
    group = parser_wan.add_mutually_exclusive_group()
    group.add_argument('--connect', action="store_true", help="Bring up connection")
    group.add_argument('--disconnect', action="store_true", help="Bring down connection")
    parser_wan.add_argument('--append-dns', dest="resolv", action="store_true",
                            help="Add the providers DNS servers to /etc/resolv.conf")
    parser_wan.add_argument('--context', type=int, help="Specify context number (default: all contexts)")

    parser_sms = sub.add_parser('sms', help="Send sms message")
    parser_sms.add_argument('--message', '-m', help="The message, if left out your editor will be opened")
    parser_sms.add_argument('destination', help="Destination number for the message")
    sub.add_parser('sms-list', help="List stored SMS messages")

    args = parser.parse_args()

    if args.action is None or args.action == "list":
        action_list()
        return

    if args.action == "poweron":
        action_power('Powered', True, 'poweron')
        return

    if args.action == "poweroff":
        action_power('Powered', False, 'poweroff')
        return

    if args.action == "online":
        action_power('Online', True, 'online')
        return

    if args.action == "offline":
        action_power('Online', False, 'offline')
        return

    if args.action == "operators":
        action_scan_operators()
        return

    if args.action == "wan":
        action_wan(args.connect, args.disconnect, args.resolv, args.context)
        return

    if args.action == "sms":
        action_sms(args.destination, args.message)
        return

    if args.action == "sms-list":
        action_sms_get()
        return

if __name__ == '__main__':
    main()
