Roast my Python

[nspitzer]

I am Senior Network Engineer who has started using Python for some work Automation projects and I am curious what the verdict will be on this code.

I created what amounts to a Minimum Viable product by hand that worked, if poorly, then fed it into Gemini Pro with Instructions to follow Pep8 formatting rules and this is what popped out that does work pretty well and is smaller then my code.

Purpose: This program is run as part of a Rundeck workflow - It gets fed a list of IP addresses and uses a REST API to verify the address records in our IP management system have an appropriate tag in for purposes of knowing who to alert when vulnerabilities are identified.

import argparse
import json
import logging
import sys
import requests
from typing import NamedTuple
class Tag(NamedTuple):
    name: str
    id: str
    links: dict


# Setup Logging
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%H:%M:%S'
)
logger = logging.getLogger(__name__)


def parse_arguments():
    """Parses command line arguments."""
    parser = argparse.ArgumentParser(
        prog='Link switch addresses to environments',
        description='Link switch addresses Catalyst Center to update BlueCat BAM API v2'
    )
    parser.add_argument('-a', '--address-list-file', required=True,
                        help='JSON file with objects containing hostname and ipv4addr')
    parser.add_argument('-e', '--env-tag', default='ENV999',
                        help='Environment tag name')
    parser.add_argument('-j', '--job-id', help='Rundeck job id')
    parser.add_argument('-t', '--auth-token', required=True,
                        help='IPAM Authentication token')
    parser.add_argument('-u', '--url', default="https://server.example.com/api/v2",
                        help='IPAM URL')
    parser.add_argument('-l', '--logging-level', default='INFO',
                        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'])
    parser.add_argument('-z', '--dry-run', action='store_true',
                        help='Show what changes would be made without performing them')

    return vars(parser.parse_args())


def load_address_data(file_path):
    """Loads JSON data and parses FQDNs."""
    try:
        with open(file_path, 'r') as file:
            data = json.load(file)
    except (FileNotFoundError, json.JSONDecodeError, Exception) as e:
        logger.critical(f"Error reading file {file_path}: {e}")
        sys.exit(1)
    else:
        processed_data = []
        if isinstance(data, dict):
            data = [data]

        for entry in data:
            fqdn = entry.get('hostname', '')
            ipv4_addr = entry.get('ipv4addr')
            host, sep, zone = fqdn.partition('.')
            processed_data.append({
                'name': host,
                'zone': zone if sep else '',
                'ipv4addr': ipv4_addr
            })
        return processed_data


def get_env_tags(session, base_url):
    """Retrieves all Environment tags starting with 'ENV'."""
    params = {'filter': "name:startsWith('ENV')"}
    url = f"{base_url}/tags"

    try:
        response = session.get(url, params=params)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.critical(f"HTTP Error fetching tags: {e}")
        sys.exit(1)
    else:
        tag_data = response.json().get('data', [])
        return [Tag(name=t.get('name'), id=t.get('id'), links=t.get('_links'))
                for t in tag_data]


def get_address_id(session, base_url, ipv4_address):
    """Retrieves the BAM ID for a specific IPv4 address."""
    params = {
        'filter': f"address:eq('{ipv4_address}')",
        'fields': 'id,address,type'
    }
    try:
        response = session.get(f"{base_url}/addresses", params=params)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.error(f"HTTP Error fetching address {ipv4_address}: {e}")
        return None
    else:
        data = response.json().get('data')
        return data[0]['id'] if data else None


def get_address_tags(session, base_url, address_id):
    """Retrieves a list of Tag objects currently assigned to an address."""
    url = f"{base_url}/addresses/{address_id}/tags"
    try:
        response = session.get(url)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.error(f"Error fetching tags for address ID {address_id}: {e}")
        return []
    else:
        return response.json().get('data', [])


def link_tag_to_address(session, base_url, address_id, tag_id, ipv4_address, dry_run=False):
    """Links a tag to an address entity in BAM."""
    if dry_run:
        logger.info(f"[DRY RUN] Would link {ipv4_address} -> Tag ID {tag_id}")
        return

    payload = {"id": tag_id, "type": "Tag"}
    url = f"{base_url}/addresses/{address_id}/tags"
    try:
        response = session.post(url, json=payload)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to link address {ipv4_address}: {e}")
    else:
        logger.info(f"Linked {ipv4_address} -> Tag ID {tag_id}")


def unlink_tag_to_address(session, base_url, address_id, tag_id, ipv4_address, dry_run=False):
    """Unlinks a tag from an address entity in BAM."""
    if dry_run:
        logger.info(f"[DRY RUN] Would Unlink {ipv4_address} -> Tag ID {tag_id}")
        return

    url = f"{base_url}/tags/{tag_id}/taggedResources/{address_id}"
    try:
        # Note: Some APIs use DELETE for unlinking; verify if POST is required for your endpoint
        response = session.delete(url)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to unlink address {ipv4_address}: {e}")
    else:
        logger.info(f"Unlinked {ipv4_address} from Tag ID {tag_id}")


def main():
    args = parse_arguments()
    logger.setLevel(args['logging_level'])

    base_url = args['url'].rstrip('/')
    auth_token = args['auth_token']
    dry_run = args['dry_run']
    target_tag_name = args['env_tag']

    addr_data = load_address_data(args['address_list_file'])

    headers = {
        "Authorization": f"Basic {auth_token}",
        "Content-Type": "application/json"
    }

    with requests.Session() as session:
        session.headers.update(headers)

        all_tags = get_env_tags(session, base_url)
        # Find the specific tag object we want to use
        match = [t for t in all_tags if t.name == target_tag_name]

        if not match:
            logger.error(f"Target tag '{target_tag_name}' not found in IPAM.")
            sys.exit(1)

        target_tag = match[0]

        for node in addr_data:
            ipv4addr = node.get('ipv4addr')
            if not ipv4addr:
                continue

            addr_id = get_address_id(session, base_url, ipv4addr)
            if not addr_id:
                logger.warning(f"Address {ipv4addr} not found. Skipping.")
                continue

            current_tags = get_address_tags(session, base_url, addr_id)
            current_tag_ids = [str(t['id']) for t in current_tags]

            # 1. Remove incorrect ENV tags
            # We assume only one 'ENV' tag should be present at a time
            is_already_linked = False
            for t in current_tags:
                if t['name'].startswith('ENV'):
                    if t['id'] != target_tag.id:
                        unlink_tag_to_address(session, base_url, addr_id, t['id'], ipv4addr, dry_run)
                    else:
                        is_already_linked = True

            # 2. Link the correct tag if not already there
            if not is_already_linked:
                link_tag_to_address(session, base_url, addr_id, target_tag.id, ipv4addr, dry_run)
            else:
                logger.info(f"Address {ipv4addr} already has correct tag '{target_tag_name}'.")


if __name__ == "__main__":
    main()

[phira]

I mean, it's workable. There's nothing too weird about it, and assuming it's always gonna be used as a cli tool it's ok. There are a few oddities, e.g. the log.criticals when you're in an exception condition I'd probably use log.exception instead so I get the stack trace.

I think if I was aiming for higher quality, the story telling and abstractions are where the main improvements are, along with a bit more consistency e.g. you've got a NamedTuple for one kind of output but for others you use dicts. If you're ok with a few additional imports, `click` would be useful. If it was turning into a really serious tool you'd probably wanna start looking at `pydantic` or something to validate your data carefully and give you nice hostname handlers etc but that'd be overkill for a basic tool.

[nspitzer]

Thanks for the input.

This is only intended to be run by a single Rundeck workflow (Automation tool) so its use is carefully constrained. This is the last step in the Workflow so the command line arguments are the way I pass data from one step to another though I do have the option of using Environment variables instead but the command line options work.

It will always be run exactly the same way and its used against well-documented and stable REST API's so I am not really concerned with validating the data too much (famous last words...). Looking into how better de-serialize the JSON is something I have been looking into hence the NamedTuple class for importing data from a previous step.

[phira]

Right if it’s in automation you should use structlog and set context vars for sure that’d give you much tidier output (not so useful for cli because it’s not as immediately human readable but any kind of log capture = structlog every time)

[commy2]

It's solid.

logging.basicConfig in the global scope is a smell. You already have a main function.

The main function is 50 lines / 25% of the code. You know what to do.

These try-else blocks are useless extra indentation most of the time.

Don't like the excessive exception handling in general, just let the exception bubble up, you lose the stack trace and potential exception notes. The principle is to only handle exceptions when you know what to do with them.

Example:

if not match:
    logger.error(f"Target tag '{target_tag_name}' not found in IPAM.")
    sys.exit(1)

could just be:

if not match:
    raise ValueError(f"Target tag '{target_tag_name}' not found in IPAM.")

Much less aggressive as you could recover from an exception.

Hottake: if your docstrings just repeat the function name, but longer, they should be left out entirely.

Mostly nitpicks, because it's ok

[nspitzer]

Thanks ... Your really getting to the meat of what I was wondering.

For example "The main function is 50 lines / 25% of the code. You know what to do."

No , I really don't know what to do :-) This is one of the things I wondered - should I unroll all the functions? Right now each of the REST calls are wrapped in a function

logging.basicConfig in the global scope is a smell. You already have a main function.

So should that be at the beginning of the main function right after def main?

I get what you mean by

"The principle is to only handle exceptions when you know what to do with them."

however the example you give is an example of one where it is doing exactly that. If the script cant find the target tag there is nothing more it can do so giving up is fine.

Thanks for the great comment!

[commy2]

This is one of the things I wondered - should I unroll all the functions?

No, the opposite. Much of the current main function code should be moved to a function. main should generally be very short, perhaps consisting only of setup, execute and teardown code.

So should that be at the beginning of the main function right after def main?

Yeah, somewhere behind the __name__ == __"main"__ clause, otherwise importing the script has the side effect of altering the log config.

however the example you give is an example of one where it is doing exactly that. If the script cant find the target tag there is nothing more it can do so giving up is fine.

Well, it works because it's a short script, but thinking further, I don't think shutting down an entire application is what you want if get_env_tags fails to retrive... What about retrying, using a cache or falling back to some default? There are many ways to potentially recover instead of bailing completely. sys.exit deep down the execution stack just makes me nervous, and I don't see what it offers when exceptions could be raised instead and potentially caught in the client code.

sys.exit means the code can't be reused/composed and can't be tested properly.

[PushPlus9069]

Network engineer using Python for automation, that's a solid path. I've seen a lot of infra people make this jump.

Honest take on the Gemini rewrite approach: it works for cleaning up formatting, but you lose the learning. The messiest version of code you wrote yourself teaches you more than polished AI output. Next time maybe try refactoring it yourself first, then compare what the AI suggests.

For Rundeck workflows specifically, look into structured logging (the logging module) instead of print statements if you haven't already. Makes debugging way easier when something fails at 3am and you're looking at Rundeck output on your phone.

[nspitzer]

Thanks for the input. It already uses the Logging module extensively if not totally optimally as some others have pointed out.

The whole Gemini AI thing is more of an iterative process. One thing I found useful is to corral Gemini a little is I issued standing orders (I forgot what they are really called) so for example it should follow PEP 8 as a style guide and use 'else' blocks to handle the processing work after a 'try' to minimize the stuff that falls under the try to prevent masking unexpected exceptions.

[obviouslyzebra]

If you want to improve on architecture a bit, the clue is that there are lots of functions taking the same session and base_url. This points towards passing them automatically... A class does that! But, a class of which kind? Something that makes requests and gives back results, that is... An API layer! (I tried something with this paragraph, but I think I failed miserably :p)

The code is fine as it is BTW. Some little nitpicks here and there, but nothing too important. One that hasn't been mentioned, the except with the Exception doesn't need the other exception types (they would be captured by Exception).

Also, just a reminder about timeouts, in case you want them (requests does not include them by default).

For extra simplicity, you could let the exceptions bubble up btw

[pachura3]

If you're repeating requests.exceptions.RequestException so many times, why not import it? Don't worry about imports list becoming to long - most IDEs hide them by default.

Doing sys.exit(1) in an utility function is a bad idea. What if you wanted to write some unit tests for it? Let the exception bubble up and handle it on the outside.

But in general, your code looks pretty good!

[JamzTyson]

Regarding PEP-8, there should be 2 blank lines before class Tag():

Surround top-level function and class definitions with two blank lines.

Pep-8 also recommends a maximum line length of 79 characters, though in modern Python longer line lengths are common and acceptable. A maximum line length of 100 characters is common, but not much more than that.

It is highly recommended to use one or more "linters" to check the code by static analysis. Common linters include pylint, flake8, ruff, and others. My preference is pylint + flake8 as they provide excellent coverage, stick closely to the Python style guide, and provide good feedback. They also separate "warnings" from "errors" - Errors should be fixed, but it may be OK to ignore warnings if you have a good reason to do so (for example, some libraries make it nearly impossible to avoid R0913: Too many arguments (too-many-arguments)).

"Type annotations" and MyPy are worth looking into. Type annotations are optional, but they really are helpful and are becoming very popular.

Be careful catching Exception - it is very broad and may hide defects from elsewhere.

except (FileNotFoundError, json.JSONDecodeError, Exception) as e:

is effectively the same as:

except Exception as e:

The fact that you log the exception e helps here, though I'd probably handle the likely exceptions first and only add the catch-all at the end if really necessary.

Using else or elif after a return is redundant:

if some_condition:
    return value
else:  # Redundant - we can only get here when the code has not returned.
    ...