En este notebook, exploraremos cómo utilizar CodeShield, una herramienta esencial para identificar vulnerabilidades de seguridad en tu código. Aprenderás a instalarlo, configurarlo y aplicarlo en diversos escenarios, desde ejemplos simples hasta integraciones con modelos de lenguaje grandes (LLM) como OpenAI y modelos alojados externamente. Al final de esta lección, serás capaz de incorporar CodeShield en tu flujo de trabajo de desarrollo para mejorar la seguridad de tus aplicaciones.

Este notebook muestra ejemplos de cómo usar CodeShield. Para más información, consulta el README del repositorio principal aquí.

Primeros Pasos

Puedes instalar CodeShield a través de PyPi usando pip, o instalarlo localmente desde el código fuente. Ambas opciones te permitirán empezar a usar la herramienta rápidamente.

Opción de Instalación 1. Instalar el paquete CodeShield. Ejecuta lo siguiente en tu terminal

Esta opción es la más sencilla y recomendada para la mayoría de los usuarios, ya que instala CodeShield directamente desde el repositorio oficial de PyPi.

pip3 install codeshield

Opción de Instalación 2. Instalar el paquete CodeShield localmente desde el código fuente

Si necesitas modificar el código fuente de CodeShield o contribuir a su desarrollo, esta opción te permite clonar el repositorio y trabajar con una instalación local.

git clone https://github.com/meta-llama/PurpleLlama
cd PurpleLlama/CodeShield
pip install .

Ejemplos

Antes de sumergirnos en los ejemplos específicos, definiremos una función auxiliar que nos permitirá ejecutar el escaneo de CodeShield y procesar su salida de manera consistente. Esto simplificará la demostración de los diferentes casos de uso.

from codeshield.cs import CodeShield

async def scan_llm_output(llm_output_code):
    result = await CodeShield.scan_code(llm_output_code)
    if result.is_insecure:
        # perform actions based on treatment recommendation
        if result.recommended_treatment == "block":
            llm_output_code = "*** Code Security issues found, blocking the code ***"
        if result.recommended_treatment == "warn":
            llm_output_code = llm_output_code + "*** Warning: The generated snippit contains insecure code ***"
    
    
    summary = "Security issue detected" if result.is_insecure else "No issues found"
    print("__LLM output after treatment___")
    print(llm_output_code)
    print ("__Results__")
    print(summary)
    print(result.recommended_treatment)
    print ("__Details__")
    print(result.issues_found)

Ejemplo 1 - Ejemplo simple con código fijo

En este primer ejemplo, utilizaremos un fragmento de código predefinido que contiene una vulnerabilidad conocida: el uso de la función de hash débil MD5. CodeShield detectará esta debilidad y nos mostrará cómo reporta tales hallazgos.

Este código simula una situación común donde se utiliza una función criptográfica obsoleta, lo cual es un riesgo de seguridad.

llm_output_code = """
def hashString(input):
    return hashlib.md5(input)
"""

await scan_llm_output(llm_output_code)

Después de ejecutar el escaneo, CodeShield analizará el fragmento de código y generará un informe detallado sobre la vulnerabilidad detectada. La salida mostrará la ubicación y la naturaleza del problema.

__LLM output after treatment___

def hashString(input):
    return hashlib.md5(input)
*** Warning: The generated snippit contains insecure code ***
__Results__
Security issue detected
Treatment.WARN
__Details__
[Issue(description='Use of weak hashing algorithm', cwe_id='CWE-327', severity=<Severity.WARNING: 'warning'>, rule='\\.getMd5Digest\\(\\)|\\.md5\\(|\\.md5Hex\\(|\\.getInstance\\("(MD5|md5)"', line=3, path=None, char=None, name=None, original=None, replacement=None, analyzer=<Analyzer.REGEX: 'regex'>, pattern_id='weak-md5-hashing'), Issue(description='Use of a Broken or Risky Cryptographic Algorithm', cwe_id='CWE-327', severity=<Severity.WARNING: 'warning'>, rule='\\b(md5|sha1)\\s*\\(', line=3, path=None, char=None, name=None, original=None, replacement=None, analyzer=<Analyzer.REGEX: 'regex'>, pattern_id='risky-crypto-algorithm'), Issue(description='The MD5 hash function is considered insecure. Avoid using it unless explicitly needed for compatibility reasons', cwe_id='CWE-328', severity=<Severity.WARNING: 'warning'>, rule='\\bhashlib\\.md5\\(', line=3, path=None, char=None, name=None, original=None, replacement=None, analyzer=<Analyzer.REGEX: 'regex'>, pattern_id='insecure-md5-hash-usage')]

Ejemplo 2 - Usar la API de OpenAI

CodeShield puede integrarse con modelos de lenguaje grandes (LLM) como los de OpenAI para mejorar sus capacidades de análisis. Para este ejemplo, necesitarás tener el paquete openai instalado y una clave de API válida.

Asegúrate de haber ejecutado pip install openai y de haber configurado tu clave de API de OpenAI como una variable de entorno o directamente en el código (para propósitos de demostración).

prompt = "please generate some example code to demonstrate strcpy usage"

import openai
client = openai.OpenAI(api_key="YOUR_OPEN_AI_KEY")
response = client.chat.completions.create(
    model= "gpt-3.5-turbo",
    messages=[
        {"role": "user", "content": prompt},
    ],
    max_tokens=1000,
)

await scan_llm_output(response.choices[0].message.content)

Ejemplo 3 - Usar un LLM alojado externamente

Además de OpenAI, CodeShield también puede trabajar con otros LLM alojados externamente. Este ejemplo demuestra cómo integrarse con modelos disponibles a través del paquete llama-recipes.

Para este ejemplo, es necesario instalar el paquete llama-recipes. Esto te permitirá conectar CodeShield a modelos de lenguaje como Llama 2, que pueden estar ejecutándose en tu entorno local o en un servidor remoto.

import os
import getpass

from llama_cookbook.inference.llm import TOGETHER, OPENAI, ANYSCALE

if "EXTERNALLY_HOSTED_LLM_TOKEN" not in os.environ:
    os.environ["EXTERNALLY_HOSTED_LLM_TOKEN"] = getpass.getpass(prompt="Provide token for LLM provider")

# Delete as appropriate
model = TOGETHER("togethercomputer/CodeLlama-13b-Instruct", os.environ["EXTERNALLY_HOSTED_LLM_TOKEN"])
model = OPENAI("gpt-4",os.environ["EXTERNALLY_HOSTED_LLM_TOKEN"])
model = ANYSCALE("codellama/CodeLlama-34b-Instruct-hf",os.environ["EXTERNALLY_HOSTED_LLM_TOKEN"])

llm_output_code = model.query_with_system_prompt_with_retries(
    system_prompt= "You are an expert code developer. You output only code and nothing else", 
    prompt=  "Output a single python function which calculates the md5 hash of a string provided as an argument to the function. Output only the code and nothing else."
    )
await scan_llm_output(llm_output_code)

Resumen

• CodeShield es una herramienta para escanear código en busca de vulnerabilidades de seguridad.
• Se puede instalar fácilmente a través de pip o desde el código fuente.
• Permite analizar código fijo para detectar patrones de seguridad débiles, como el uso de MD5.
• Se integra con APIs de LLM como OpenAI para análisis avanzados de código.
• También soporta la conexión a LLM alojados externamente, como los disponibles a través de llama-recipes.