CGenPy.cs

/*
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See the LICENSE file in the project root for full license information.
 */

using System.IO;
using System.Diagnostics;
using System.Collections.Generic;
using System.Linq;

namespace CodeGen
{
    /// <summary> Python TSS code generator </summary>
    class CGenPy : CodeGenBase
    {
        public CGenPy(string rootDir) : base(rootDir + @"src\", "TpmExtensions.py.snips") {}

        void WriteAutoGeneratedSourceHeader()
        {
            TpmNamedConstant ver = TpmTypes.LookupConstant("TPM_SPEC_VERSION");
            Write("\"\"\"\r\n" +
                    " * Copyright(c) Microsoft Corporation. All rights reserved. \r\n" +
                    " * Licensed under the MIT License. \r\n" +
                    " * See the LICENSE file in the project root for full license information. \r\n" +
                    "\"\"\"\r\n" +
                    "\r\n" +
                    "\"\"\"\r\n" +
                    " * This file is automatically generated from the TPM 2.0 rev. " +
                    (ver.NumericValue / 100.0).ToString("0.00") + " specification documents.\r\n" +
                    " * Do not edit it directly.\r\n" +
                    "\"\"\"\r\n");
        }

        internal override void Generate()
        {
            WriteAutoGeneratedSourceHeader();
            Write("from .TpmStructure import *");
            Write("from .TpmEnum import *");
            Write("");

            // First generate enums
            foreach (var e in TpmTypes.Get<TpmEnum>())
                GenEnum(e);
            foreach (var b in TpmTypes.Get<TpmBitfield>())
                GenBitfield(b);

            Write("from .Crypt import *" + "\r\n");

            // Then generate unions and structures
            GenUnions();
            foreach (var s in TpmTypes.Get<TpmStruct>())
                GenStruct(s);

            File.WriteAllText(RootDir + "TpmTypes.py", GeneratedCode.ToString());
            GeneratedCode.Clear();

            // Now generate the TPM methods
            GenCommands();
            File.WriteAllText(RootDir + "Tpm.py", GeneratedCode.ToString());
            GeneratedCode.Clear();
        }

        void GenEnum(TpmType e, List<TpmNamedConstant> elements)
        {
            // e.UnderlyingType.GetSize()
            TabIn($"class {e.Name}(TpmEnum): # {e.UnderlyingType.SpecName}");
            WriteComment(e);

            foreach (var v in elements)
            {
                // Skip TPM_RC.H enumerator (an alias hiding SUCCESS value)
                if (v.Name == "H")
                    continue;

                Write("");
                Write($"{v.Name} = {v.Value}");
                WriteComment(v);

                // Backward compat
                if (v.Name == "PW")
                {
                    WriteComment("Deprecated: use PW instead");
                    Write($"RS_PW = {v.Value}");
                }
            }
            string type = e is TpmEnum ? "enum" : "bitfield";
            TabOut($"# {type} {e.Name}");
        } // GenEnum()

        void GenEnum(TpmEnum e)
        {
            GenEnum(e, e.Members);
        }

        void GenBitfield(TpmBitfield bf)
        {
            GenEnum(bf, GetBifieldElements(bf));
        }

        /// <summary>
        /// C unions of the TPM spec are translated into Python classes implementing
        /// an interface defining the union 
        /// </summary>
        void GenUnions()
        {
            // Base class for union interfaces
            TabIn("class TpmUnion(TpmMarshaller):");
            WriteComment("TPM union interface");
            Write("@abc.abstractmethod");
            Write($"def GetUnionSelector({This}): pass  # returns TPM_ALG_ID | TPM_CAP | TPM_ST");
            TabOut("");

            // Union interfaces definitions
            var unions = TpmTypes.Get<TpmUnion>();
            foreach (TpmUnion u in unions)
            {
                TabIn($"class {u.Name}(TpmUnion):");
                WriteComment(u);
                Write("pass");
                TabOut("");
            }

            // Union factory
            TabIn("class UnionFactory:");
            Write("@staticmethod");
            TabIn("def create(unionType, selector):");
            WriteComment("Args:\n" +
                         "    unionType (string): union type name\n" +
                         "    selector (TPM_ALG_ID|TPM_CAP|TPM_ST): enum value\n" +
                         "        specifying the union member to instantiate");
            string elIf = "if";
            foreach (TpmUnion u in unions)
            {
                Write($"{elIf} unionType == '{u.Name}':");
                elIf = "elif";

                TabIn();
                foreach (UnionMember m in u.Members)
                {
                    string newObject = m.Type.IsElementary() ? TargetLang.Null : $"{m.Type.Name}()";
                    Write($"if selector == {m.SelectorValue.QualifiedName}: return {newObject}");
                }
                TabOut();
            }
            TabIn("else:");
            Write("raise(Exception('UnionFactory.create(): Unrecognized union type \"{}\"'.format(unionType)))");
            TabOut("raise(Exception('UnionFactory.create(): Unknown selector value \"{s}\" for union \"{u}\"'.format(s = str(selector), u = unionType)))");
            TabOut("# create()", false);
            TabOut("# class UnionFactory()");
        }

        void GenGetUnionSelector(TpmStruct s)
        {
            string selType = GetUnionMemberSelectorInfo(s, out string selVal);
            if (selType != null)
            {
                Write("");
                TabIn($"def GetUnionSelector({This}): # {selType}");
                WriteComment("TpmUnion method");
                Write($"return {selVal}");
                TabOut();
            }
        }

        void GenMarshalingMethod(bool dirTo, TpmStruct s)
        {
            var fields = s.MarshalFields;
            if (s.DerivedFrom != null || fields.Count() == 0)
                return;

            string  dir = dirTo ? "to" : "initFrom",
                    proto = $"def {dir}Tpm({This}, buf):";

            var marshalOps = dirTo ? GetToTpmFieldsMarshalOps(fields)
                                   : GetFromTpmFieldsMarshalOps(fields);

            Write("");
            TabIn(proto);
            WriteComment("TpmMarshaller method");
            foreach (var op in marshalOps)
                Write($"{op}");
            TabOut("", false);
        } // GenMarshalingMethod()

        /// <summary>
        /// Structures in C are classes in TypeScript
        /// </summary>
        /// <param name="s"></param>
        void GenStruct(TpmStruct s)
        {
            bool hasBase = s.DerivedFrom != null;
            string className = s.Name;
            string classBases = hasBase ? s.DerivedFrom.Name
                              : !s.IsCmdStruct() ? "TpmStructure"
                              : s.Info.IsRequest() ? "ReqStructure" : "RespStructure";

            // If this struct is not derived from another one and is a member of one or more unions,
            // it must implement the corresponding union interfaces
            if (!s.IsCmdStruct() && !hasBase)
            {
                string unionInterfaces = string.Join(", ", s.ContainingUnions.Select(u => u.Name));
                if (unionInterfaces != "")
                    classBases += ", " + unionInterfaces;
            }

            // Class header
            TabIn($"class {className} ({classBases}):");

            // Javadoc comment for the data structure
            string annotation = s.Comment + eol;
            var fields = s.FieldHolder.NonTagFields;
            if (fields.Length != 0)
            {
                annotation += eol + "Attributes:" + eol;
                foreach (var f in fields)
                    annotation += GetParamComment(f, "    ", "        ") + eol;
            }

            //
            // Field defining constructor
            //

            if (fields.Count() == 0)
            {
                TabIn($"def __init__({This}):");
                WriteComment(annotation);
                Write("pass");
                TabOut();
            }
            else
            {
                string ctorParams = string.Join(", ", fields.Select(f => f.Name + $" = {f.GetInitVal()}"));
                TabIn($"def __init__({This}, {ctorParams}):");
                WriteComment(annotation);

                // If a non-trivial base class is present, all member fields are contained there,
                // so just pass through the initialization parameters
                if (hasBase)
                {
                    string baseInitParams = string.Join(", ", fields.Select(f => f.Name));
                    //Write($"{s.DerivedFrom}.__init__({baseInitParams})");
                    Write($"super({className}, {This}).__init__({baseInitParams})");
                }
                else
                {
                    foreach (var f in fields)
                        Write($"{ThisMember}{f.Name} = {f.Name}");
                }
                TabOut();

                //
                // Named union tag getters (instead of the corresponding fields in the TPM 2.0 spec)
                //
                foreach (var sel in s.TagFields)
                {
                    if (sel.MarshalType != MarshalType.UnionSelector)
                        continue;

                    var unionField = sel.RelatedUnion;
                    var u = (TpmUnion)unionField.Type;
                    Write("");
                    Write("@property");
                    TabIn($"def {sel.Name}({This}): # {sel.TypeName}");
                    WriteComment(sel.Comment);
                    if (u.NullSelector == null)
                        Write($"return {unionField.Name}.GetUnionSelector()");
                    else
                    {
                        Write($"return {unionField.Name}.GetUnionSelector() if {unionField.Name}" +
                                                                         $" else {u.NullSelector.QualifiedName}");
                    }
                    TabOut();
                }

                foreach (var f in fields)
                {
                    if (f.MarshalType != MarshalType.UnionObject)
                        continue;

                    var u = (TpmUnion)f.Type;
                    var sel = (f as UnionField).UnionSelector;
                }
            }

            GenGetUnionSelector(s);

            //
            // Marshaling
            //
            GenMarshalingMethod(true, s);
            GenMarshalingMethod(false, s);

            Write("");
            Write($"@staticmethod");
            TabIn("def fromTpm(buf):");
            WriteComment($"Returns new {className} object constructed from its marshaled representation in the given TpmBuffer buffer");
            Write($"return buf.createObj({className})");
            TabOut("");

            Write($"@staticmethod");
            TabIn("def fromBytes(buffer):");
            WriteComment($"Returns new {className} object constructed from its marshaled representation in the given byte buffer");
            Write($"return TpmBuffer(buffer).createObj({className})");
            TabOut();

            var info = s.IsCmdStruct() ? s.Info as CmdStructInfo : null;
            if (info != null && (info.NumHandles != 0 || info.SessEncSizeLen != 0))
            {
                if (info.NumHandles != 0)
                {
                    Write("");
                    Write($"def numHandles({This}): return {info.NumHandles}");
                    if (info.IsRequest())
                    {
                        string handles = string.Join(", ", s.Fields.Take(info.NumHandles).Select(f => ThisMember + f.Name));
                        Write("");
                        Write($"def numAuthHandles({This}): return {info.NumAuthHandles}");
                        Write("");
                        Write($"def getHandles({This}): return [{handles}]");
                    }
                    else
                    {
                        Debug.Assert(info.NumHandles == 1 && info.NumAuthHandles == 0);
                        Write("");
                        Write($"def getHandle({This}): return {ThisMember}{s.Fields[0].Name}");
                        Write("");
                        Write($"def setHandle({This}, h): {ThisMember}{s.Fields[0].Name} = h");
                    }
                }
                if (info.SessEncSizeLen != 0)
                {
                    Debug.Assert(info.SessEncValLen != 0);
                    Write("");
                    Write($"def sessEncInfo({This}): return SessEncInfo({info.SessEncSizeLen}, {info.SessEncValLen})");
                }
            }

            //
            // Custom members
            //
            InsertSnip(s.Name);

            TabOut($"# {className}");
        } // GenStruct()

        void GenCommands()
        {
            WriteAutoGeneratedSourceHeader();
            Write("from .TpmBase import *" + "\r\n");

            Write($"class Tpm(TpmBase):");
            TabIn();
            foreach (var req in TpmTypes.Get<TpmStruct>().Where(s => s.Info.IsRequest()))
            {
                GenCommand(req);
            }
            TabOut();
            Write("# class Tpm");
        }

        void GenCommand(TpmStruct req)
        {
            var resp = GetRespStruct(req);
            var cmdName = GetCommandName(req);
            string respType = resp.Name;
            string cmdCode = "TPM_CC." + cmdName;
            var reqFields = req.NonTagFields;
            var respFields = resp.NonTagFields;

            if (ForceJustOneReturnParm.Contains(cmdName))
                respFields = respFields.Take(1).ToArray();

            int numOutParms = respFields.Count();
            string returnType = numOutParms == 1 ? respFields[0].TypeName
                              : numOutParms == 0 ? "void" : respType;

            // javadoc annotation
            string annotation = req.Comment;
            if (reqFields.Count() != 0)
            {
                annotation += eol + eol + "Args:";
                foreach (var f in reqFields)
                    annotation += eol + GetParamComment(f, "    ", "        ");
            }
            if (respFields.Count() != 0)
            {
                annotation += eol + eol + "Returns:";
                annotation += eol + GetReturnComment(respFields);
            }

            // method definition
            string paramList = This;
            string reqStructInitList = "";
            foreach (var f in reqFields)
            {
                paramList += ", " + f.Name;  // ": " + f.TypeName + ", ";
                reqStructInitList += (reqStructInitList.Length == 0 ? "" : ", ") + f.Name;
            }

            TabIn($"def {cmdName}({paramList}):");
            WriteComment(annotation);

            Write($"req = {req.Name}({reqStructInitList})");
            Write($"respBuf = {ThisMember}dispatchCommand({cmdCode}, req)");
            respType = numOutParms > 0 ? ", " + respType : "";
            if (numOutParms == 1)
            {
                Write($"res = {ThisMember}processResponse(respBuf{respType})");
                Write($"return res.{respFields[0].Name} if res else None");
            }
            else
                Write($"return {ThisMember}processResponse(respBuf{respType})");

            TabOut();
            Write($"# {cmdName}()");
            Write("");
        } // GenCommand()

        protected override void WriteComment(string comment, bool wrap = true)
        {
            WriteComment(comment, "\"\"\" ", "", "\n\"\"\"");
        }
    } // class CGenPy
}