Methods check password for PostgreSQL

plpython3u.sql

-- Copyright 2019-2022 Jonathan S. Katz
-- 
-- MIT License
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-- copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
-- 
-- The above copyright notice and this permission notice shall be included in all
-- copies or substantial portions of the Software.
-- 
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-- SOFTWARE.

-- Generate the password hashes / verifiers for use in PostgreSQL

CREATE OR REPLACE FUNCTION pyw_usr_cheсk_pwd(
	curpwd character varying,
	hashed_password character varying)
    RETURNS boolean
    LANGUAGE 'plpython3u'
    COST 10
    STABLE PARALLEL SAFE 
AS $BODY$
	import base64
	import hashlib
	import hmac
	import stringprep
	import unicodedata
	
	hp = hashed_password.split("$")[1].split(":")
	iterations = int(hp[0])
	salt = base64.b64decode(hp[1])
	password = curpwd
	
	def bytes_xor(a, b):
		"""XOR two bytestrings together"""
		return bytes(a_i ^ b_i for a_i, b_i in zip(a, b))	

	def normalize_password(password):
		"""Normalize the password using PostgreSQL-flavored SASLprep. For reference:
		https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/common/saslprep.c
		using the `pg_saslprep` function
		Implementation borrowed from asyncpg implementation:
		https://github.com/MagicStack/asyncpg/blob/master/asyncpg/protocol/scram.pyx#L263
		"""
		SASLPREP_STEP3 = (
				stringprep.in_table_a1,  # PostgreSQL treats this as prohibited
				stringprep.in_table_c12,
				stringprep.in_table_c21_c22,
				stringprep.in_table_c3,
				stringprep.in_table_c4,
				stringprep.in_table_c5,
				stringprep.in_table_c6,
				stringprep.in_table_c7,
				stringprep.in_table_c8,
				stringprep.in_table_c9,
		)

		normalized_password = password
		# if the password is an ASCII string or fails to encode as an UTF8
		# string, we can return
		try:
				normalized_password.encode("ascii")
		except UnicodeEncodeError:
				pass
		else:
				return normalized_password
		# Step 1 of SASLPrep: Map. Per the algorithm, we map non-ascii space
		# characters to ASCII spaces (\x20 or \u0020, but we will use ' ') and
		# commonly mapped to nothing characters are removed
		# Table C.1.2 -- non-ASCII spaces
		# Table B.1 -- "Commonly mapped to nothing"
		normalized_password = u"".join(
				[' ' if stringprep.in_table_c12(c) else c
				 for c in normalized_password if not stringprep.in_table_b1(c)])
		# If at this point the password is empty, PostgreSQL uses the original
		# password
		if not normalized_password:
				return password
		# Step 2 of SASLPrep: Normalize. Normalize the password using the
		# Unicode normalization algorithm to NFKC form
		normalized_password = unicodedata.normalize(
				'NFKC', normalized_password)
		# If the password is not empty, PostgreSQL uses the original password
		if not normalized_password:
				return password
		# Step 3 of SASLPrep: Prohobited characters. If PostgreSQL detects any
		# of the prohibited characters in SASLPrep, it will use the original
		# password
		# We also include "unassigned code points" in the prohibited character
		# category as PostgreSQL does the same
		for c in normalized_password:
				if any([in_prohibited_table(c) for in_prohibited_table in SASLPREP_STEP3]):
						return password
		# Step 4 of SASLPrep: Bi-directional characters. PostgreSQL follows the
		# rules for bi-directional characters laid on in RFC3454 Sec. 6 which
		# are:
		# 1. Characters in RFC 3454 Sec 5.8 are prohibited (C.8)
		# 2. If a string contains a RandALCat character, it cannot containy any
		#    LCat character
		# 3. If the string contains any RandALCat character, an RandALCat
		#    character must be the first and last character of the string
		# RandALCat characters are found in table D.1, whereas LCat are in D.2
		if any([stringprep.in_table_d1(c) for c in normalized_password]):
				# if the first character or the last character are not in D.1,
				# return the original password
				if not (stringprep.in_table_d1(normalized_password[0]) and
								stringprep.in_table_d1(normalized_password[-1])):
						return password
				# if any characters are in D.2, use the original password
				if any([stringprep.in_table_d2(c) for c in normalized_password]):
						return password
		# return the normalized password
		return normalized_password

	def scram_sha_256_generate_salted_password(password, iterations, salt, digest):
		"""This follows the "Hi" algorithm specified in RFC5802"""
		# first, need to normalize the password using PostgreSQL-flavored SASLprep
		normalized_password = normalize_password(password)
		# convert the password to a binary string - UTF8 is safe for SASL (though there are SASLPrep rules)
		p = normalized_password.encode("utf8")
		# the initial signature is the salt with a terminator of a 32-bit string ending in 1
		ui = hmac.new(p, salt + b'\x00\x00\x00\x01', digest)
		# grab the initial digest
		u = ui.digest()
		# for X number of iterations, recompute the HMAC signature against the password
		# and the latest iteration of the hash, and XOR it with the previous version
		for x in range(iterations - 1):
				ui = hmac.new(p, ui.digest(), hashlib.sha256)
				# this is a fancy way of XORing two byte strings together
				u = bytes_xor(u, ui.digest())
		return u

	def encrypt(password, salt, iterations):
		salted_password = scram_sha_256_generate_salted_password(
				password, iterations, salt, hashlib.sha256)
		client_key = hmac.HMAC(salted_password, b"Client Key", hashlib.sha256)
		stored_key = hashlib.sha256(client_key.digest()).digest()
		server_key = hmac.HMAC(salted_password, b"Server Key", hashlib.sha256)
		encrypted_password = "SCRAM-SHA-256".encode("utf-8") + b"$" + \
				("{}".format(iterations)).encode("utf-8") + b":" + \
				base64.b64encode(salt) + b"$" + \
				base64.b64encode(stored_key) + b":" + \
				base64.b64encode(server_key.digest())
		return encrypted_password

	return encrypt(password, salt, iterations).decode('utf-8') == hashed_password
	
$BODY$;

CREATE OR REPLACE FUNCTION usr_cheсk_pwd(
	usr_login character varying,
	usr_pwd character varying)
    RETURNS boolean
    LANGUAGE 'plpgsql'
    COST 100
    STABLE SECURITY DEFINER PARALLEL SAFE 
    SET search_path=bpd
AS $BODY$
DECLARE
	hashed_password character varying;
	result boolean default false;
BEGIN
	SELECT rolpassword INTO hashed_password FROM pg_authid WHERE rolname = usr_login;
		
	IF NOT(hashed_password IS NULL) THEN
		result = bpd.pyw_usr_cheсk_pwd( usr_pwd, hashed_password);
	END IF;
	RETURN result;
END;
$BODY$;