Useful Python Snippets

Logging

Configure rich logging

Install dependencies

pip install rich

import logging.config

logging.config.dictConfig(
    {
        'version': 1,
        'disable_existing_loggers': False,
        'formatters': {'basic': {'style': '{', 'format': '{message}'}},
        'handlers': {
            'rich': {
                '()': 'rich.logging.RichHandler',
                'omit_repeated_times': False,
                'highlighter': None,
            }
        },
        'loggers': {__name__: {'level': 'INFO', 'handlers': ['rich']}},
    }
)

Local Network

Dependencies

• psutil

  • psutil.net_if_addrs()

  Return the addresses associated to each NIC (network interface card) installed on the system as a dictionary whose keys are the NIC names and value is a list of named tuples for each address assigned to the NIC

• gufo-ping

  • Example from docs

    import asyncio
    import sys
    
    from gufo.ping import Ping
    
    
    async def main(addr: str) -> None:
        ping = Ping()
        r = await ping.ping(addr)
        print(r)
    
    
    if __name__ == "__main__":
        asyncio.run(main(sys.argv[1]))
    

Example

scan.py

Install dependencies

pip install rich gufo-ping psutil

scan.py

#!/usr/bin/env python3

import asyncio
import ipaddress
import socket
from collections import OrderedDict
from typing import Dict

import psutil
from gufo.ping import Ping
from psutil._common import snicaddr
from rich.console import Console
from rich.highlighter import NullHighlighter
from rich.progress import Progress, SpinnerColumn, TextColumn

console = Console(highlighter=NullHighlighter())

print = console.print

def get_self_ips(family: socket.AddressFamily = socket.AddressFamily.AF_INET) -> Dict[str, snicaddr]:
    return {
        name: addr
        for name, addrs in psutil.net_if_addrs().items()
        for addr in addrs
        if addr.family == family
    }


def subnet_mask_to_cidr(subnet_mask: str) -> int:
    # Convert subnet mask to binary string
    binary_str = ''.join([bin(int(x) + 256)[3:] for x in subnet_mask.split('.')])
    # Count the number of '1's in the binary string to get the CIDR notation
    return binary_str.count('1')


def print_ips(family: socket.AddressFamily = socket.AddressFamily.AF_INET):
    addrs = get_self_ips(family)
    for name, address in addrs.items():
        print(f'{name.ljust(15)} {address}')


async def ping_scan(interface: str):
    net_info = get_self_ips()
    try:
        addr: snicaddr = net_info[interface]
    except KeyError as e:
        print(f'[bold red]{type(e).__name__}[/]: Invalid interface: {interface}')
        return

    cidr = subnet_mask_to_cidr(addr.netmask)
    network = ipaddress.ip_network(f'{addr.address}/{cidr}', strict=False)
    ips = list(map(str, network.hosts()))

    progress_bar = Progress(
        SpinnerColumn(), TextColumn('{task.description}'),
        console=console, transient=True
    )
    with progress_bar as prog:
        prog.add_task(f'Pinging {len(ips)} IPs...')
        ping = Ping()
        async with asyncio.TaskGroup() as tg:
            tasks = {ip: tg.create_task(ping.ping(ip)) for ip in ips}

    return OrderedDict({
        ip: r * 1000
        for ip, task in tasks.items()
        if (r := task.result()) is not None
    })


if __name__ == '__main__':
    for name, addr in get_self_ips().items():
        print(f'[bold green]{name.ljust(18)}[/]{addr.address}')

    interface = 'eth0'
    pings = asyncio.run(ping_scan(interface))

    print(f'  [yellow]{interface}[/]  '.center(40, '-'))
    for ip, p in pings.items():
        print(f'[bold blue]{ip.ljust(15)}[/]{p:< 6.1f} ms')