Command-line python script to complete crazyhouse FEN by giving black all unused pieces in hand

zhume_fen.py

# === FENify tsumehouse ===
# Takes a crazyhouse FEN containing just board position and command-line 
# argument for white hand pieces, and outputs correct FEN where black has all 
# the remaining pieces in hand.
# 
# --- Example ---
# Input: "4r3/5k2/5pp1/5N2/8/8/8/8[] w - - 0 1" "BQ"
# Output: "4r3/5k2/5pp1/5N2/8/8/8/8[QBqrrrbbbnnnpppppppppppppp] w - - 0 1"
# (this is a 3te crazyhouse tsume)

import argparse
import operator
import subprocess

FULL_PIECE_COUNT = [2, 4, 4, 4, 16] # queen, rook, bishop, knight, pawn
PIECES = {"Q" : 0, "R" : 1, "B" : 2, "N" : 3, "P" : 4}
PIECES_FROM_IDX = ["Q", "R", "B", "N", "P"]

def count_fen_pieces(fen):
    piece_count = [0, 0, 0, 0, 0]
    position = fen.split("[")[0]
    for counter, char in enumerate(position):
        if char.isalpha():
            type = char.upper()
            if type in PIECES.keys():
                # check if piece is a promoted pawn
                if position[counter+1] == "~":
                    piece_count[PIECES["P"]] += 1
                else:
                    piece_count[PIECES[type]] += 1
    return piece_count;

def count_hand_pieces(hand):
    piece_count = [0, 0, 0, 0, 0]
    for char in hand:
        if char.upper() in PIECES.keys():
            piece_count[PIECES[char.upper()]] += 1
    return piece_count;

def generate_hand_string(piece_count):
    hand_list = []
    for counter, number in enumerate(piece_count):
        hand_list.extend(PIECES_FROM_IDX[counter] * number)
    return "".join(hand_list)

# !!! SHELL INJECTION VULNERABILITY
def copy_to_clipboard(text):
    cmd = "echo " + text.strip() + "|clip"
    return subprocess.check_call(cmd, shell=True)
    

# "main" function
parser = argparse.ArgumentParser(description="Generate correct tsumehouse FEN from diagram and specified white (semekata) hand pieces.")
parser.add_argument("input_fen", type=str, help="Input FEN containing board position and empty pockets for both sides. Assumes [] delimiters for pockets.")
parser.add_argument("hand_pieces", type=str, default="", help="String specifying white (semekata) hand pieces. Capital piece symbols, English notation (QRBNP), no spaces; e.g. 'RRN' or 'RNR' for 2 rooks and 1 knight.")
parser.add_argument("-c", "--clipboard", action="store_true", help="Automatically copy the output to the clipboard.")
args = parser.parse_args()

num_board_pieces = count_fen_pieces(args.input_fen)
num_white_pieces = count_hand_pieces(args.hand_pieces)
num_black_pieces = map(operator.sub, FULL_PIECE_COUNT, num_board_pieces)
num_black_pieces = map(operator.sub, num_black_pieces, num_white_pieces)
#assert num_black_pieces all entries positive or zero
white_string = generate_hand_string(num_white_pieces).upper()
black_string = generate_hand_string(num_black_pieces).lower()

fen = args.input_fen.split("[")
output_fen = fen[0] + "[" + white_string + black_string + fen[1]
print(output_fen)

if args.clipboard:
    copy_to_clipboard(output_fen)