#!/usr/bin/env python
#
# pg_corrupt
#
# Read in a byte of a PostgreSQL relation (a table or index) and allow writing
# it back with an altered value.
#
# Greg Smith <greg@2ndquadrant.com>
# Copyright (c) 2013, Heroku, Inc.
# Released under The PostgreSQL Licence
#

from os.path import join
from subprocess import Popen,PIPE
import sys
import psycopg2

class Operations:
    SHOW=0
    WRITE=1
    AND=2
    OR=3
    XOR=4
    text=['show','write','and','or','xor']

def controldata_blocks_per_segment(pgdata):
    blocks_per_seg = 131072
    try:
        # TODO This doesn't work when called in an Emacs compile shell
        out, err = Popen("pg_controldata %s" % pgdata, stdout=PIPE, shell=True).communicate()
        control_data=out.splitlines()
        for c in control_data:
            if c.startswith("Blocks per segment of large relation:"):
                blocks_per_seg=int(c.split(":")[1])
    except:
        print "Cannot determine blocks per segment, using default of",blocks_per_seg
    return blocks_per_seg

def get_table_info(conn,relation):
    cur = conn.cursor()
    q="SELECT \
       current_setting('data_directory') AS data_directory, \
       pg_relation_filepath(oid), \
       current_setting('block_size') AS block_size, \
       pg_relation_size(oid), \
       relpages \
       FROM pg_class \
       WHERE relname='%s'" % relation
    cur.execute(q)
    if cur.rowcount != 1:
        print "Error: did not return 1 row from pg_class lookup of %s" % relation
        return None
    table_info={}

    for i in cur:
        table_info['relation']=relation
        table_info['pgdata'] = i[0]
        table_info['filepath'] = i[1]
        table_info['block_size'] = int(i[2])
        table_info['relation_size'] = i[3]
        table_info['relpages'] = i[4]
        table_info['base_file_name']=join(i[0],i[1])
        table_info['blocks_per_seg'] = controldata_blocks_per_segment(table_info['pgdata'])
        table_info['bytes_per_seg'] = table_info['block_size'] * table_info['blocks_per_seg']

    cur.close()
    return table_info

def operate(table_info,byte_offset,operation,value):
    if byte_offset > table_info['relation_size']:
        print "Error:  trying to change byte %s but relation %s is only %s bytes" % \
            (byte_offset, table_info['relation'],table_info['relation_size'])
        return

    if byte_offset < 0:
        print "Error:  cannot use negative byte offsets"
        return

    file_name=table_info['base_file_name']
    file_seq=int(byte_offset / table_info['bytes_per_seg'])
    if file_seq > 0:
        file_name=file_name + ".%s" % file_seq
        file_offset = byte_offset - file_seq * table_info['bytes_per_seg']
    else:
        file_offset = byte_offset
    
    print "Reading byte",file_offset,"within file",file_name
    f = open(file_name, mode='r+b')
    f.seek(file_offset)
    current=f.read(1)
    current_int=ord(current)
    print "Current byte=",current_int,"/ $%s" % current.encode('hex')

    if operation==Operations.WRITE:
        out_int=value
    elif operation==Operations.AND:
        out_int=current_int & value
    elif operation==Operations.OR:
        out_int=current_int | value
    elif operation==Operations.XOR:
        out_int=current_int ^ value
    elif operation==Operations.SHOW:
        return None
    else:
        print "Unsupported operation type"
        return None

    out=chr(out_int)
    print "Modified byte=",out_int,"/ $%s" % out.encode('hex')

    f.seek(file_offset)
    f.write(out)
    f.close()
    print "File modified successfully"

def usage():
    print "pg_corrupt relation operation [offset value]"
    print "Operation is one of",Operations.text
    print "All operations except for show require an offset and value"

def main():
    # TODO Replace this with a proper optparse setup
    argc=len(sys.argv)
    if argc<3 or argc==4:
        usage()
        sys.exit(1)

    # Initialize optional parameters
    byte_offset=0
    value=0

    relation=sys.argv[1]
    operation=Operations.text.index(sys.argv[2])

    if operation!=Operations.SHOW:
        if argc>4:
            value=int(sys.argv[4])
            byte_offset=int(sys.argv[3])
        else:
            usage()
            sys.exit(1)

    conn = psycopg2.connect("")
    info=get_table_info(conn,relation)
    conn.close()
    if not info is None:
        operate(info,byte_offset,operation,value)

if __name__=="__main__":
    main()