Adding Extensions

Adding Extensions

Relevant source files

Purpose and Scope

This document provides technical guidance for extending WSHawk's capabilities through three primary extension mechanisms: vulnerability detection modules, mutation strategies, and payload collections. Each extension point follows a specific integration pattern that ensures compatibility with the scanner's orchestration engine and reporting system.

For setting up a development environment before adding extensions, see Development Environment Setup. For understanding the overall codebase structure, see Project Structure and Architecture.

Sources: CONTRIBUTING.md L92-L114

Extension Points Overview

WSHawk provides three distinct extension mechanisms, each serving a specific role in the testing pipeline:

| Extension Type | Location | Purpose | Integration Method | | --- | --- | --- | --- | | Vulnerability Detection Modules | wshawk/ | Add new attack vector testing capabilities | Import and call from WSHawkV2 | | Mutation Strategies | wshawk/mutators/ | Implement WAF bypass techniques | Inherit BaseMutator, register in __init__.py | | Payload Collections | payloads/ | Expand attack vector databases | Add .txt or .json files, update WSPayloads |

Sources: CONTRIBUTING.md L92-L114

setup.py L16-L62

scanner_v2.py L28-L680

Adding Vulnerability Detection Modules

Module Structure

Vulnerability detection modules integrate with the WSHawkV2 scanner class by implementing test methods that follow the established pattern. Each module must:

  1. Accept a WebSocket connection object
  2. Return a list of vulnerability dictionaries
  3. Use VulnerabilityVerifier for result validation
  4. Log findings using the Logger class
  5. Include CVSS scoring metadata

Sources: scanner_v2.py L143-L213

scanner_v2.py L28-L76

CONTRIBUTING.md L94-L101

Implementation Pattern

The standard test method pattern follows this structure, as demonstrated in existing modules:

Method Signature:

Core Components:

  1. Initialization - Load payloads and configure test parameters
  2. Intelligence Integration - Use MessageIntelligence.inject_payload_into_message() for context-aware injection
  3. Payload Execution - Send test vectors via WebSocket connection
  4. Response Analysis - Use VulnerabilityVerifier.verify_<vuln_type>() for validation
  5. Result Collection - Append findings to self.vulnerabilities with required fields

Required Vulnerability Dictionary Fields:

| Field | Type | Description | Example | | --- | --- | --- | --- | | type | str | Vulnerability classification | "SQL Injection" | | severity | str | CVSS severity level | "HIGH", "CRITICAL" | | confidence | str | Detection confidence | "HIGH", "MEDIUM", "LOW" | | description | str | Technical description | "SQL error in response" | | payload | str | Test vector used | "' OR '1'='1" | | response_snippet | str | Evidence (first 200 chars) | response[:200] | | recommendation | str | Remediation guidance | "Use parameterized queries" |

Example Implementation Pattern (from SQL injection module):

Sources: scanner_v2.py L143-L213

scanner_v2.py L215-L293

Scanner Integration

To integrate a new vulnerability module with the scanner:

Step 1: Add test method to WSHawkV2 class in scanner_v2.py

Step 2: Call test method in run_intelligent_scan() at scanner_v2.py L545-L680

Step 3: Optionally add to CLI interfaces:

Sources: scanner_v2.py L545-L680

CONTRIBUTING.md L94-L101

Adding Mutation Strategies

Mutator Architecture

Mutation strategies implement WAF bypass techniques by transforming payloads before injection. The system uses a base class pattern for consistent integration.

Sources: CONTRIBUTING.md L102-L108

wshawk/mutators/init.py

(implied structure)

Implementation Requirements

Step 1: Create mutator file in wshawk/mutators/ directory:

wshawk/mutators/custom_mutator.py

Step 2: Inherit from BaseMutator:

Step 3: Register mutator in wshawk/mutators/__init__.py:

Sources: CONTRIBUTING.md L102-L108

Mutation Strategy Examples

Common mutation patterns for WAF bypass:

| Strategy | Input | Output Example | Use Case | | --- | --- | --- | --- | | Case Variation | SELECT * FROM | sElEcT * FrOm | Case-sensitive filters | | Encoding | <script> | %3Cscript%3E, \x3Cscript\x3E | Character blacklists | | Comment Injection | ' OR 1=1 | '/**/OR/**/1=1 | Whitespace filters | | Unicode Normalization | <script> | <script> (fullwidth) | Unicode-aware WAFs | | Concatenation | SELECT | SEL'+'ECT | String concatenation bypass |

Sources: Architecture diagram context from [Diagram 4](https://github.com/noobforanonymous/wshawk/blob/0eebedf4/Diagram 4)

and CONTRIBUTING.md L102-L108

Adding Payload Collections

Payload File Formats

WSHawk supports two payload file formats, both loaded via package data configuration:

Text Format (.txt):

  • One payload per line
  • UTF-8 encoding required
  • Comments not supported
  • Direct loading via WSPayloads class methods

JSON Format (.json):

  • Structured payload metadata
  • Category-based organization
  • Context information per payload
  • Used for complex vulnerability types

Sources: setup.py L50-L55

CONTRIBUTING.md L109-L114

Text Format Implementation

Step 1: Create payload file in payloads/ directory:

payloads/custom_vectors.txt

Step 2: Add payloads (one per line, no comments):

' OR '1'='1
' OR '1'='1'--
' OR '1'='1'/*
admin'--
admin'#

Step 3: Update WSPayloads class in wshawk/main.py

:

Sources: setup.py L50-L55

CONTRIBUTING.md L109-L114

JSON Format Implementation

Step 1: Create JSON payload file:

payloads/custom/advanced_vectors.json

Step 2: Define structured payload format:

Step 3: Add loader method to WSPayloads:

Sources: setup.py L50-L55

Package Data Configuration

Ensure payload files are included in package distribution via setup.py L49-L55

:

This configuration ensures payload files are bundled with:

  • PyPI package distribution
  • Docker image builds
  • Local development installations

Sources: setup.py L49-L55

Integration with Scanner Intelligence

MessageIntelligence Integration

New vulnerability modules should leverage MessageIntelligence for context-aware payload injection:

Usage Pattern in Test Methods:

Sources: scanner_v2.py L87-L141

scanner_v2.py L163-L171

ServerFingerprinter Integration

Use ServerFingerprinter to load database/language-specific payloads:

This pattern is demonstrated in scanner_v2.py L153-L159

and scanner_v2.py L305-L311

Sources: scanner_v2.py L153-L159

scanner_v2.py L305-L311

VulnerabilityVerifier Integration

All test methods must use VulnerabilityVerifier for result validation:

Available verifier methods:

  • verify_sql_injection(response, payload)
  • verify_xss(response, payload)
  • verify_command_injection(response, payload)
  • verify_path_traversal(response, payload)

Sources: scanner_v2.py L185-L201

scanner_v2.py L244-L282

Testing and Validation

Development Testing

Before submitting extensions, validate functionality:

Step 1: Run quick module tests:

Step 2: Test against local WebSocket server:

Step 3: Verify output artifacts:

  • Check wshawk_report_*.html for proper vulnerability formatting
  • Verify CVSS scores are present
  • Confirm evidence snippets are captured

Sources: CONTRIBUTING.md L52-L61

Code Quality Requirements

Extensions must meet project standards:

| Requirement | Description | Tool/Method | | --- | --- | --- | | PEP 8 Compliance | Follow Python style guidelines | flake8, black | | Type Hints | Add type annotations to functions | mypy | | Docstrings | Document all classes and methods | Google/NumPy style | | Error Handling | Wrap risky operations in try/except | N/A | | Rate Limiting | Respect rate_limiter.acquire() | Use await asyncio.sleep() |

Example Docstring Format:

Sources: CONTRIBUTING.md L43-L49

Submission Process

Pull Request Checklist

Before submitting extensions:

  • All tests pass (python tests/test_modules_quick.py)
  • Code follows PEP 8 style guidelines
  • Type hints added to new functions
  • Docstrings added to all classes/methods
  • Error handling implemented for network operations
  • Integration point updated (scanner, mutator registry, or WSPayloads)
  • Example usage documented in docstrings
  • CVSS scoring included for vulnerability modules
  • Rate limiting respected in test methods

Documentation Requirements

Include in pull request description:

  1. Purpose: What vulnerability/technique does this extension detect/implement?
  2. Integration: How does it integrate with existing scanner components?
  3. Testing: What testing was performed? Include example output.
  4. Dependencies: Any new dependencies required?
  5. Breaking Changes: Does this modify existing interfaces?

Sources: CONTRIBUTING.md L85-L91

Review Process

Extension submissions undergo:

  1. Automated CI/CD Checks - Build verification and test execution
  2. Code Review - Maintainer evaluation of implementation quality
  3. Security Review - Assessment of potential security implications
  4. Integration Testing - Verification of scanner compatibility

Expect 3-7 days for initial review feedback.

Sources: CONTRIBUTING.md L70-L91

Diagram 6 from architecture overview