FTL/test/api/libs/FTLAPI.py

#!/bin/python3
# Pi-hole: A black hole for Internet advertisements
# (c) 2023 Pi-hole, LLC (https://pi-hole.net)
# Network-wide ad blocking via your own hardware.
#
# FTL Engine - auxiliary files
# API test script
#
# This file is copyright under the latest version of the EUPL.
# Please see LICENSE file for your rights under this license.

from enum import Enum
import random
import requests
from typing import List
import json
from hashlib import sha256

url = "http://pi.hole/api/auth"

class AuthenticationMethods(Enum):
	RANDOM = 0
	HEADER = 1
	BODY = 2
	COOKIE = 3

# Class to query the FTL API
class FTLAPI():

	auth_method = "?"

	def __init__(self, api_url: str, password: str = None):
		self.api_url = api_url
		self.endpoints = {
			"get": [],
			"post": [],
			"put": [],
			"patch": [],
			"delete": []
		}
		self.errors = []
		self.session = None
		self.verbose = False

		# Login to FTL API
		if password is not None:
			self.login(password)
			if self.session is None or 'valid' not in self.session or not self.session['valid']:
				raise Exception("Could not login to FTL API")

	def login(self, password: str = None):
		# Check if we even need to login
		response = self.GET("/api/auth")

		# Check if we are already logged in or authentication is not
		# required
		if response is None:
			raise Exception("No response from FTL API")
		if 'session' not in response:
			raise Exception("FTL returned invalid challenge item")
		if 'session' in response and response['session']['valid'] == True:
			self.session = response["session"]
			print(response)
			if password is not None:
				raise Exception("Password provided but API does not require authentication")
			return

		response = self.POST("/api/auth", {"password": password})
		if "error" in response:
			raise Exception("FTL returned error: " + json.dumps(response["error"]))
		if 'session' not in response:
			raise Exception("FTL returned invalid response item")
		self.session = response["session"]


	def get_jsondata_headers_cookies(self, authenticate: AuthenticationMethods):
		# Add session ID to the request (if any)
		json_data = None
		headers = None
		cookies = None
		if self.session is not None and 'sid' in self.session:
			# Pick a random authentication method if requested
			# Try again if the method comes out as random again
			while authenticate == AuthenticationMethods.RANDOM:
				authenticate = random.choice(list(AuthenticationMethods))

			# Add the session ID to the request
			if authenticate == AuthenticationMethods.HEADER:
				headers = {"X-FTL-SID": self.session['sid']}
			elif authenticate == AuthenticationMethods.BODY:
				json_data = {"sid": self.session['sid'] }
			elif authenticate == AuthenticationMethods.COOKIE:
				# Cookie authentication needs both the session ID and the CSRF header
				cookies = {"sid": self.session['sid'] }
				headers = { "X-CSRF-Token": self.session['csrf'] }

			self.auth_method = authenticate.name

		return json_data, headers, cookies


	# Query the FTL API (GET) and return the response
	def GET(self, uri: str, params: List[str] = [], expected_mimetype: str = "application/json", authenticate: AuthenticationMethods = AuthenticationMethods.BODY):
		self.errors = []
		try:
			# Add parameters to the URI (if any)
			if len(params) > 0:
				uri = uri + "?" + "&".join(params)

			# Get json_data, headers and cookies
			json_data, headers, cookies = self.get_jsondata_headers_cookies(authenticate)

			if self.verbose:
				print("GET " + self.api_url + uri + " with json_data: " + json.dumps(json_data))

			# Query the API
			with requests.get(url = self.api_url + uri, json = json_data, headers=headers, cookies=cookies) as response:
				if self.verbose:
					print(json.dumps(response.json(), indent=4))
				if expected_mimetype == "application/json":
					return response.json()
				else:
					return response.content
		except Exception as e:
			self.errors.append("Exception when GETing from FTL: " + str(e))
			return None


	# Query the FTL API (POST) and return the response
	def POST(self, uri: str, json_data: dict = {}, authenticate: AuthenticationMethods = AuthenticationMethods.HEADER, files = None):
		self.errors = []
		try:
			# Get json_data, headers and cookies
			_, headers, cookies = self.get_jsondata_headers_cookies(authenticate)

			if self.verbose:
				print("POST " + self.api_url + uri + " with json_data: " + json.dumps(json_data))

			# Query the API
			with requests.post(url = self.api_url + uri, json = json_data, files = files, headers=headers, cookies=cookies) as response:
				if self.verbose:
					print(json.dumps(response.json(), indent=4))
				return response.json()
		except Exception as e:
			self.errors.append("Exception when POSTing to FTL: " + str(e))
			return None


	# Query the endpoints from FTL for comparison with the OpenAPI specs
	def get_endpoints(self):
		try:
			# Get all endpoints from FTL and sort them for comparison
			response = self.GET("/api/endpoints")
			for method in response["endpoints"]:
				for endpoint in response["endpoints"][method]:
					self.endpoints[method].append(endpoint["uri"] + endpoint["parameters"])
			for method in self.endpoints:
				self.endpoints[method] = sorted(self.endpoints[method])
		except Exception as e:
			print("Exception when pre-processing endpoints from FTL: " + str(e))
			exit(1)

		return self.endpoints