# -*- coding: utf-8 -*-
import base64




def safe_base64_decode(s):
    # pass
    if type(s) == bytes:
        s = str(s, encoding='utf-8')
    # 题目中的 b'YWJjZA==' 是一个bytes类型，于是直接用utf-8转换为string以便后面添加等号

    if len(s) % 4 == 0:
        return base64.b64decode(s)    
    else:
        return base64.b64decode( s + '='*(4 - len(s) % 4) )




# 测试:
assert b'abcd' == safe_base64_decode(b'YWJjZA=='), safe_base64_decode('YWJjZA==')
assert b'abcd' == safe_base64_decode(b'YWJjZA'), safe_base64_decode('YWJjZA')
print('ok')