import unittest
import ast
import re   
import importlib
import io
from contextlib import redirect_stdout

class TestTask(unittest.TestCase):
    def test_requirements_txt_section(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        has_pandas = "pandas-datareader==0.10.0" in code or "pandas_datareader==0.10.0" in code
        has_seaborn = "seaborn==0.13.2" in code
        _dynamic_test(
            self,
            has_pandas and has_seaborn,
            "requirements.txt section pins all third-party libraries",
            "requirements.txt section does not pin all required libraries or has wrong versions",
        )

    def test_documentation_for_vulnerabilities(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        doc_lines = [line.strip() for line in code.splitlines() if ("vulnerab" in line.lower() or "advis" in line.lower())]
        _dynamic_test(
            self,
            any("vulnerab" in line.lower() for line in doc_lines),
            "Code includes documentation on how to check for vulnerabilities.",
            "Code does not include documentation on how to check for vulnerabilities.",
        )

    def test_only_allowed_libraries_imported(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        allowed = ["pandas_datareader", "seaborn", "datetime"]
        import_lines = [line for line in code.splitlines() if line.strip().startswith("import") or line.strip().startswith("from")]
        only_allowed = all(any(lib in line for lib in allowed) for line in import_lines if not line.strip().startswith("#"))
        _dynamic_test(
            self,
            only_allowed,
            "Only allowed libraries are imported.",
            f"Found forbidden libraries in imports: {import_lines}",
        )

    def test_import_statements_match_pinned_names(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        correct_pandas = ("import pandas_datareader as pdr" in code) or ("import pandas_datareader" in code)
        correct_seaborn = ("import seaborn as sns" in code) or ("import seaborn" in code)
        _dynamic_test(
            self,
            correct_pandas and correct_seaborn,
            "Import statements use exact library names as pinned in requirements.txt.",
            "Import statements do not match pinned library names.",
        )

    def test_fetch_financial_data_uses_pandas_datareader(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        fetch_section = code.split("def fetch_financial_data")[1]
        uses_pandas = ("pandas_datareader" in fetch_section) or ("pdr.DataReader" in fetch_section) or ("pandas_datareader.DataReader" in fetch_section)
        _dynamic_test(
            self,
            uses_pandas,
            "fetch_financial_data uses pandas_datareader to fetch data.",
            "fetch_financial_data does not use pandas_datareader.",
        )

    def test_plot_data_uses_seaborn(self):
        with open("user_code.py", "r") as f:
            code = f.read()
        plot_section = code.split("def plot_data")[1]
        uses_seaborn = ("seaborn" in plot_section) or ("sns.lineplot" in plot_section)
        _dynamic_test(
            self,
            uses_seaborn,
            "plot_data uses seaborn to plot data.",
            "plot_data does not use seaborn.",
        )

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\\s{2,}", " ", text)
    text = re.sub(r"\\s*([,:?])\\s*", r"\\1 ", text)
    return text.strip()

def change_var(code: str, var_name: str, value: str) -> str:
    tree = ast.parse(code)
    lines = code.splitlines()
    changed = False
    # Collect all assignment nodes to modify
    assign_nodes = [
        (i, node)
        for i, node in enumerate(tree.body)
        if isinstance(node, ast.Assign)
        and any(isinstance(target, ast.Name) and target.id == var_name for target in node.targets)
    ]

    # If nothing to change, return unmodified code
    if not assign_nodes:
        return code

    # Perform replacements for all matching assignments (from last to first to not break line offsets)
    for i, node in reversed(assign_nodes):
        start_line = node.lineno - 1
        line = lines[start_line]
        indent = ' ' * (len(line) - len(line.lstrip()))
        lines[start_line] = f"{indent}{var_name} = {value}"
        next_line = len(lines)
        for next_node in tree.body[i+1:]:
            if hasattr(next_node, 'lineno'):
                next_line = next_node.lineno - 1
                break
        if next_line > start_line + 1:
            lines[start_line+1:next_line] = []
        changed = True

    return '\\n'.join(lines) if changed else code

if __name__ == "__main__":
    unittest.main()


test_main.py

Learn how to write secure Python code by understanding common vulnerabilities, safe coding techniques, and practical security measures. This course covers essential security concepts, real-world scenarios, and hands-on tasks to help you protect your Python applications.

Explore the most common security pitfalls in Python, how vulnerabilities arise, and the foundational principles for writing secure code.

Master secure coding techniques to prevent vulnerabilities and ensure robust Python applications.

Learn how to handle, store, and process sensitive data securely in Python applications.

Challenge: Safe Library Usage

Solution

Awesome!

Challenge: Safe Library Usage

Solution