From 5eb1c2e6e40589f1c4f06b3859ae11686cd834c7 Mon Sep 17 00:00:00 2001 From: Andre Sookram Date: Wed, 4 Sep 2024 18:39:55 -0400 Subject: [PATCH] Update README.md and ffufai.py to support Ollama --- README.md | 5 ++++- ffufai.py | 23 ++++++++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index a73a702..f4eaee1 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,10 @@ ffufai is an AI-powered wrapper for the popular web fuzzer ffuf. It automaticall ``` export ANTHROPIC_API_KEY='your-api-key-here' ``` - + Or for Ollama: + ``` + export OLLAMA_API_KEY='ollama' + ``` You can add these lines to your `~/.bashrc` or `~/.zshrc` file to make them permanent. ## Usage diff --git a/ffufai.py b/ffufai.py index c7ef118..30f8070 100755 --- a/ffufai.py +++ b/ffufai.py @@ -9,15 +9,19 @@ import anthropic from urllib.parse import urlparse + def get_api_key(): + ollama_key = os.getenv('OLLAMA_API_KEY') openai_key = os.getenv('OPENAI_API_KEY') anthropic_key = os.getenv('ANTHROPIC_API_KEY') - if anthropic_key: + if ollama_key: + return ('ollama', ollama_key) + elif anthropic_key: return ('anthropic', anthropic_key) elif openai_key: return ('openai', openai_key) else: - raise ValueError("No API key found. Please set OPENAI_API_KEY or ANTHROPIC_API_KEY.") + raise ValueError("No API key found. Please set OLLAMA_API_KEY or OPENAI_API_KEY or ANTHROPIC_API_KEY.") def get_headers(url): try: @@ -52,7 +56,20 @@ def get_ai_extensions(url, headers, api_type, api_key, max_extensions): JSON Response: """ - if api_type == 'openai': + if api_type == 'ollama': + client = OpenAI( + base_url = 'http://127.0.0.1:11434/v1', #set to IP address to network server if applicable + api_key=api_key, # required, but unused + ) + response = client.chat.completions.create( + model="llama3.1", #set local model details + messages=[ + {"role": "system", "content": "You are a helpful assistant that suggests file extensions for fuzzing based on URL and headers."}, + {"role": "user", "content": prompt} + ] + ) + return json.loads(response.choices[0].message.content.strip()) + elif api_type == 'openai': client = OpenAI(api_key=api_key) response = client.chat.completions.create( model="gpt-4o",