How CyberStrikeAI Implements Role-Based Testing: A Deep Dive into the Go Source Code

CyberStrikeAI isolates security testing scenarios by assigning roles that bundle a user prompt, tool whitelist, and skill list, automatically enforcing these constraints at runtime when processing AI agent requests.

CyberStrikeAI is an open-source security testing platform that leverages role-based testing to transform general-purpose AI models into specialized penetration-testing profiles. By configuring distinct roles with specific tools and skills, security teams can ensure that the AI operates within strictly defined boundaries during vulnerability assessments. This article examines the Go source code implementation in the Ed1s0nZ/CyberStrikeAI repository to reveal how role definitions are structured, persisted, and enforced during request processing.

The Role Configuration Schema

At the core of the implementation lies the RoleConfig struct defined in internal/config/config.go. This structure defines the shape of a testing role, encapsulating all parameters necessary to constrain the AI's behavior.

type RoleConfig struct {
    Name        string
    Description string
    UserPrompt  string
    Icon        string
    Tools       []string
    MCPs        []string
    Skills      []string
    Enabled     bool
}

The UserPrompt field injects context-specific instructions, while Tools and Skills arrays explicitly whitelist the utilities and capabilities available to the AI during that session. When Enabled is set to false, the role becomes unavailable for selection without deleting its configuration.

REST API for Role Lifecycle Management

The RoleHandler in internal/handler/role.go exposes full CRUD operations through REST endpoints including GET /api/roles, POST /api/roles, PUT /api/roles/:name, and DELETE /api/roles/:name. When creating or updating a role, the system normalizes role names and persists configurations to individual YAML files within the roles/ directory.

The saveConfig method handles filesystem persistence by sanitizing filenames and marshaling the configuration:

// sanitizeFileName removes unsafe characters
filename := sanitizeFileName(role.Name)
filepath := filepath.Join(h.config.RolesDir, filename+".yaml")
data, _ := yaml.Marshal(role)
os.WriteFile(filepath, data, 0644)

During updates, the UpdateRole function manages key renames by deleting stale files and writing new ones, ensuring the in-memory config.Roles map remains synchronized with the disk state.

Runtime Role Enforcement in the Agent Loop

When processing a testing request, the AgentHandler in internal/handler/agent.go checks for a role parameter in the incoming request. If specified and not set to "默认" (default), the handler retrieves the corresponding RoleConfig and modifies the execution context.

The following logic demonstrates how the system enforces role-based constraints:

if req.Role != "" && req.Role != "默认" {
    role := h.config.Roles[req.Role]
    finalMessage = role.UserPrompt + "\n\n" + req.Message
    roleTools  = role.Tools   // Whitelisted tools only
    roleSkills = role.Skills  // Specific skills only
}
result, err := h.agent.AgentLoopWithProgress(..., roleTools, roleSkills)

Clients invoke this behavior by posting to /api/agent-loop with a JSON payload specifying the target role:

POST /api/agent-loop HTTP/1.1
Content-Type: application/json

{
  "message": "Scan the target example.com for open ports and vulnerable endpoints.",
  "role": "WebAppPenTest"
}

This ensures the AI operates exclusively with the tools and prompt context defined for that specific penetration-testing profile.

Role Support for Corporate Messenger Bots

Beyond the primary web API, CyberStrikeAI extends role-based testing to corporate messenger integrations through the RobotHandler in internal/handler/robot.go. This handler supports platforms such as DingTalk and Lark, applying the same role enrichment logic to bot conversations.

The implementation mirrors the agent handler's approach, using h.config.Roles to prepend the role's UserPrompt and restrict the available toolset. This allows security teams to initiate constrained testing scenarios directly from chat interfaces while maintaining the same isolation guarantees as the REST API.

Skill-to-Role Reverse Mapping

The SkillsHandler in internal/handler/skills.go provides a reverse lookup capability through the getRolesBoundToSkill method. This functionality enables the frontend to display which testing profiles utilize a specific skill, aiding in dependency management and configuration audits.

When accessing GET /api/skills/:skillName/bound_roles, the handler iterates through all configured roles:

GET /api/skills/sql-injection-testing/bound_roles HTTP/1.1

The response identifies all roles referencing that skill:

{
  "skill": "sql-injection-testing",
  "bound_roles": ["WebAppPenTest", "API PenTest"],
  "bound_count": 2
}

This reverse mapping ensures administrators understand the impact of modifying or deleting specific skills across the role-based testing infrastructure.

Summary

• RoleConfig struct in internal/config/config.go defines the schema for role-based testing, bundling prompts, tools, and skills.
• RoleHandler in internal/handler/role.go manages CRUD operations and persists roles as individual YAML files in the roles/ directory.
• AgentHandler enforces runtime constraints by prepending role prompts and restricting tool invocation to whitelisted sets.
• RobotHandler extends these capabilities to corporate messenger bots like DingTalk and Lark.
• SkillsHandler provides reverse lookups to identify which roles depend on specific skills.

Frequently Asked Questions

How does CyberStrikeAI handle the default role when no specific role is selected?

The system treats requests with an empty role string or the value "默认" as default sessions, skipping role-specific prompt injection and tool restrictions. This allows general-purpose AI behavior when specialized testing profiles are not required.

Where does CyberStrikeAI store role configurations on the filesystem?

Role configurations are stored as individual YAML files within the directory specified by config.RolesDir, defaulting to roles/. Each role file uses a sanitized version of the role name as its filename, created via the saveConfig method in internal/handler/role.go.

Can a role restrict which external tools the AI is permitted to invoke?

Yes, the Tools and Skills arrays in the RoleConfig struct explicitly whitelist permitted utilities. When processing a request, AgentHandler passes these arrays to AgentLoopWithProgress, ensuring the AI can only access the designated testing tools during that session.

What happens to the filesystem when a role is renamed?

The UpdateRole function in internal/handler/role.go handles renaming by deleting the old YAML file and creating a new one with the sanitized updated name. This ensures filesystem consistency with the in-memory config.Roles map while preventing orphaned configuration files.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}