Unit 1: Python for defenders · Lesson 4 of 6 · 20 min

Functions and modules

Wrap a check in a function, reuse it from another script and keep it safe to import.

mission

You are
the engineer turning one-off checks into team tools
Situation
Three scripts on the team parse the inventory slightly differently, and each one has its own bug. You will write the parser once and share it.
Goal
Build an inventory module with a parser and a risk check, then import it from a report script.
Hand in
A report of hosts whose risk score is 2 or more.

A module as a tool rack

01

Write the module

A module is just a .py file. Functions take input, return output and print nothing, so any script can reuse them.

learner@py-lab
cat > scripts/inventory_lib.py <<'EOF'
RISKY = {21: "ftp", 23: "telnet", 3389: "remote desktop", 445: "smb"}


def parse(line):
    name, role, system, ports = line.split()
    return {"name": name, "role": role, "os": system, "ports": [int(p) for p in ports.split(",")]}


def load(path="/home/learner/data/inventory.txt"):
    return [parse(line) for line in open(path) if line.strip()]


def risk(host):
    return sorted(RISKY[p] for p in host["ports"] if p in RISKY)


if __name__ == "__main__":
    for host in load():
        print(host["name"], risk(host))
EOF
02

Run the module directly

The if __name__ == "__main__" block runs only when you execute the file, not when another script imports it.

learner@py-lab
python3 scripts/inventory_lib.py
web01 []
web02 ['remote desktop']
db01 ['telnet']
db02 []
mail01 []
vpn01 ['telnet']
...
03

Import it from a report

learner@py-lab
cat > scripts/risk_report.py <<'EOF'
from inventory_lib import load, risk

hosts = load()
scored = [(len(risk(h)), h["name"], risk(h)) for h in hosts]
for score, name, found in sorted(scored, reverse=True):
    if score >= 2:
        print(f"{score}  {name:<10} {', '.join(found)}")
EOF
cd scripts && python3 risk_report.py
3  legacy01   ftp, smb, telnet
3  dc02       remote desktop, smb, telnet
2  file03     remote desktop, smb
2  file02     smb, telnet
2  dc03       remote desktop, smb
2  dc01       remote desktop, smb

you should see Only hosts with two or more risky services, worst first.

defender tip Python finds inventory_lib because it sits next to the script you run. Run from another folder and the import fails, which is why real tools become packages.

04

Test a function in isolation

Small functions can be checked by hand in the REPL with made-up input.

>>
import sys; sys.path.insert(0, 'scripts')
from inventory_lib import parse, risk
risk(parse('test01 web ubuntu-24.04 22,23,445'))
['smb', 'telnet']
risk(parse('test02 web ubuntu-24.04 22,443'))
[]

you should see A host with telnet and smb, then a clean host with an empty list.

What a defender takes away

  1. 1Write each check once, as a function that returns data.
  2. 2Keep printing in the script, logic in the module.
  3. 3Guard script-only code with if __name__ == "__main__".

Questions

0 of 3 solved

Answer each one from what you find on your lab machine.

  1. Show hint

    The first line of the report.

  2. Show hint

    Pipe the report into wc -l.

  3. Show hint

    sum("smb" in risk(h) for h in load())