Update .gitmodules files in bulk

update_gitmodules.mjs

// update_gitmodules.mjs
import fs from 'fs/promises';
import path from 'path';
import { fileURLToPath } from 'url';
import readline from 'readline';
import { parse, stringify } from 'ini';
import simpleGit from 'simple-git';
import { createPatch } from 'diff';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const githubOrg = 'balena-os';
const repoYamlMatch = 'yocto-based OS image';
const prBranch = 'kyle/update-gitmodules';
const submodulesDir = path.join(__dirname, 'submodules');
const cacheDir = path.join(__dirname, '.cache');

const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});

function prompt(question) {
  return new Promise((resolve) => {
    rl.question(question, resolve);
  });
}

function delay(ms) {
  return new Promise(resolve => setTimeout(resolve, ms));
}

async function checkRepoYaml(octokit, owner, repo) {
  try {
    const { data } = await octokit.repos.getContent({
      owner,
      repo,
      path: 'repo.yml',
    });
    const content = Buffer.from(data.content, 'base64').toString('utf-8');
    return content.includes(repoYamlMatch);
  } catch (error) {
    if (error.status === 404) {
      console.log(`${repo}: repo.yml not found, skipping...`);
      return false;
    }
    throw error;
  }
}

async function checkGitModules(octokit, owner, repo) {
  try {
    const { data } = await octokit.repos.getContent({
      owner,
      repo,
      path: '.gitmodules',
    });
    const content = Buffer.from(data.content, 'base64').toString('utf-8');
    
    const gitmodules = parse(content);

    for (const section in gitmodules) {
      if (section.startsWith('submodule ')) {
        const submodule = gitmodules[section];
        if (!submodule.branch) {
          return true;
        }
      }
    }

    return false;
  } catch (error) {
    if (error.status === 404) {
      return false;
    }
    throw error;
  }
}

async function cloneOrUpdateSubmodule(submoduleUrl, submodulePath) {
  const localPath = path.join(submodulesDir, submodulePath);
  const git = simpleGit();

  try {
    await fs.access(localPath);
    // Submodule exists, update it
    const submoduleGit = simpleGit(localPath);
    await submoduleGit.fetch(['--all']);
  } catch (error) {
    if (error.code === 'ENOENT') {
      // Submodule doesn't exist, clone it
      try {
        await fs.mkdir(path.dirname(localPath), { recursive: true });
        await git.clone(submoduleUrl, localPath);
      } catch (error) {
        console.error(`Error cloning submodule: ${error.message}`);
        throw error;
      }
    } else {
      console.error(`Error accessing submodule: ${error.message}`);
      throw error;
    }
  }
  return localPath;
}

async function findBranchesContainingCommit(submodulePath, commitSha) {
  const git = simpleGit(submodulePath);
  const branches = await git.branch(['--remotes', '--contains', commitSha]);
  return branches.all
    .map(branch => branch.replace('origin/', ''))
    .filter(branch => branch && !branch.includes('HEAD'));
}

async function getSubmoduleBranch(logRepo, submoduleUrl, submodulePath, submoduleCommit, currentGitmodules) {

  switch (submodulePath) {
    case 'layers/meta-balena':
    case 'balena-yocto-scripts':
    case 'contracts':
    case 'tests/autohat':
      return 'master';
  }

  try {
    const localSubmodulePath = await cloneOrUpdateSubmodule(submoduleUrl, submodulePath);

    console.log(`${logRepo}: ${submodulePath}: Getting branches containing commit ${submoduleCommit}...`);
    const matchingBranches = await findBranchesContainingCommit(localSubmodulePath, submoduleCommit);

    if (matchingBranches.length > 0) {

      const currentPokyBranch = currentGitmodules['submodule "layers/poky"'].branch || undefined;
      const currentMetaOEBranch = currentGitmodules['submodule "layers/meta-openembedded"'].branch || undefined;

      // console.log(`${logRepo}: Current poky branch: ${currentPokyBranch}`);
      // console.log(`${logRepo}: Current meta-oe branch: ${currentMetaOEBranch}`);

      if (matchingBranches.includes(currentPokyBranch)) {
        console.log(`${logRepo}: ${submodulePath}: Using current poky branch: ${currentPokyBranch}`);
        return currentGitmodules['submodule "layers/poky"']?.branch;
      }

      if (matchingBranches.includes(currentMetaOEBranch)) {
        console.log(`${logRepo}: ${submodulePath}: Using current meta-oe branch: ${currentMetaOEBranch}`);
        return currentGitmodules['submodule "layers/meta-openembedded"']?.branch;
      }

      console.log(`${logRepo}: ${submodulePath}: Matching branches:`);
      console.log(`0) Skip this submodule`);
      matchingBranches.forEach((branch, index) => {
        console.log(`${index + 1}) ${branch}`);
      });

      let choice = await prompt('Enter your choice: ');
      let choiceIndex = parseInt(choice);

      while (isNaN(choiceIndex) || choiceIndex < 0 || choiceIndex > matchingBranches.length) {
        console.log('Invalid choice. Please enter a number.');
        choice = await prompt('Enter your choice: ');
        choiceIndex = parseInt(choice);
      }

      switch (choiceIndex) {
        case 0:
          return null;
        default:
          return matchingBranches[choiceIndex-1];
      }

    }

    console.log(`${logRepo}: ${submodulePath}: Unable to determine branch!`);
    return null;

  } catch (error) {
    console.error(`Error getting branches for ${submodulePath}: ${error.message}`);
    return null;
  }
}

async function createOrResetBranch(octokit, owner, repo, branchName, baseBranch = 'master') {
  console.log(`${repo}: Creating or resetting branch ${branchName}...`);

  try {
    // Get the SHA of the latest commit on the base branch
    const { data: baseRef } = await octokit.git.getRef({
      owner,
      repo,
      ref: `heads/${baseBranch}`,
    });

    const baseSha = baseRef.object.sha;

    try {
      // Try to get the current branch (to check if it exists)
      await octokit.git.getRef({
        owner,
        repo,
        ref: `heads/${branchName}`,
      });

      // If the branch exists, update it to point to the latest commit of the base branch
      await octokit.git.updateRef({
        owner,
        repo,
        ref: `heads/${branchName}`,
        sha: baseSha,
        force: true,
      });

      console.log(`${repo}: Branch ${branchName} has been reset to ${baseBranch}`);
    } catch (error) {
      if (error.status === 404) {
        // If the branch doesn't exist, create it
        await octokit.git.createRef({
          owner,
          repo,
          ref: `refs/heads/${branchName}`,
          sha: baseSha,
        });

        console.log(`${repo}: Branch ${branchName} has been created`);
      } else {
        throw error;
      }
    }
  } catch (error) {
    console.error(`${repo}: Error creating or resetting branch ${branchName}:`, error.message);
    throw error;
  }
}

function customStringify(obj) {
  let output = stringify(obj);
  
  // Add spaces around '=' and tabs before properties
  output = output.replace(/^([^\[].+?)=(.*)$/gm, '\t$1 = $2')  // Add tab before properties and spaces around '='
             .replace(/\n+/g, '\n');  // Remove extra newlines

  return output;
}

async function updateGitmodules(octokit, owner, repo) {
  console.log(`${repo}: Checking .gitmodules for updates...`);

  try {
    const { data: gitmodulesFile } = await octokit.repos.getContent({
      owner,
      repo,
      path: '.gitmodules',
    });

    const gitmodulesContent = Buffer.from(gitmodulesFile.content, 'base64').toString('utf-8');
    const gitmodules = parse(gitmodulesContent);
    // let updatedContent = gitmodulesContent;

    let changes = false;

    for (const section in gitmodules) {
      if (section.startsWith('submodule ')) {
        const submodule = gitmodules[section];
        if (!submodule.branch) {
          const submodulePath = submodule.path;
          const submoduleUrl = submodule.url;

          const { data: submoduleCommit } = await octokit.repos.getContent({
            owner,
            repo,
            path: submodulePath,
          });

          const branch = await getSubmoduleBranch(repo, submoduleUrl, submodulePath, submoduleCommit.sha, gitmodules);
          if (branch) {
            console.log(`${repo}: ${submodulePath}: Setting branch to ${branch}`);
            gitmodules[section].branch = branch;
            changes = true;
          }
        }
      }
    }

    if (!changes) {
      console.log(`${repo}: No changes needed in .gitmodules`);
      return false;
    }

    const updatedContent = customStringify(gitmodules);
    
    console.log(`${repo}: Changes to .gitmodules:`);
    const diff = createPatch('.gitmodules', gitmodulesContent, updatedContent);
    console.log(diff);
    // console.log(updatedContent);

    const createPRChoice = await prompt(`${repo}: Push changes to branch ${prBranch}? (y/n): `);
    if (createPRChoice.toLowerCase() !== 'y') {
      console.log(`${repo}: Skipping branch/commit creation...`);
      return false;
    }

    console.log(`${repo}: Creating/resetting branch ${prBranch}...`);
    await createOrResetBranch(octokit, owner, repo, prBranch);

    const result = await octokit.repos.createOrUpdateFileContents({
      owner,
      repo,
      path: '.gitmodules',
      message: 'Update .gitmodules with submodule branch information\n\nChangelog-entry: Update .gitmodules with submodule branch information',
      content: Buffer.from(updatedContent).toString('base64'),
      sha: gitmodulesFile.sha,
      branch: prBranch,
    });

    console.log(`${repo}: File update result:`, result.status, result.data.commit.sha);

    return true;
  } catch (error) {
    console.error(`${repo}: Error updating .gitmodules:`, error.message);
    if (error.response) {
      console.error(`${repo}: Error response:`, error.response.data);
    }
    throw error;
  }
}

async function createPR(octokit, owner, repo) {

  const createPRChoice = await prompt(`Create PR for ${repo}? (y/n): `);
  if (createPRChoice.toLowerCase() !== 'y') {
    console.log(`Skipping PR creation for ${repo}`);
    return;
  }

  try {
    const { data: existingPRs } = await octokit.pulls.list({
      owner,
      repo,
      state: 'open',
      head: `${githubOrg}:${prBranch}`,
    });

    if (existingPRs.length > 0) {
      console.log(`PR already exists for ${repo}. Exiting...`);
      return;
    }
  } catch (error) {
    console.error(`Error checking existing PRs for ${repo}: ${error.message}`);
    throw error;
  }

  const { data: pr } = await octokit.pulls.create({
    owner,
    repo,
    title: 'Update .gitmodules with submodule branch information',
    head: prBranch,
    base: 'master',
    body: 'Changelog-entry: Update .gitmodules with submodule branch information',
  });

  console.log(`Created PR: ${pr.html_url}`);
}

async function processRepository(repo, octokit) {
  const repoName = repo.name;

  if (await shouldSkipRepo(repo)) {
    return;
  }

  if (repo.archived) {
    console.log(`${repoName}: Skipping archived repository...`);
    await markRepoAsSkipped(repo);
    return;
  }

  // Add a delay between each repository to avoid rate limiting
  await delay(1000);

  const isYoctoRepo = await checkRepoYaml(octokit, githubOrg, repoName);
  if (!isYoctoRepo) {
    console.log(`${repoName}: Skipping non-yocto repository...`);
    await markRepoAsSkipped(repo);
    return;
  }

  const hasUnsetModuleBranch = await checkGitModules(octokit, githubOrg, repoName);
  if (!hasUnsetModuleBranch) {
    console.log(`${repoName}: Skipping repository with all module branches set...`);
    await markRepoAsSkipped(repo);
    return;
  }

  const changes = await updateGitmodules(octokit, githubOrg, repoName);
  if (changes) {
    await createPR(octokit, githubOrg, repoName);
  }
}

async function shouldSkipRepo(repo) {
  // Check if .ignore file exists for this repo
  const ignoreFilePath = `${cacheDir}/${repo.name}.ignore`;
  try {
    await fs.access(ignoreFilePath);
    // console.log(`${repo.name}: Skipping repo with an ignore file...`);
    return true; // .ignore file exists, so skip this repo
  } catch (error) {
    return false; // .ignore file does not exist, do not skip
  }
}

async function markRepoAsSkipped(repo) {
  // Write an .ignore file for this repo
  const ignoreFilePath = `${cacheDir}/${repo.name}.ignore`;
  await fs.mkdir(path.dirname(ignoreFilePath), { recursive: true });
  await fs.writeFile(ignoreFilePath, 'Skipped due to conditions met on previous run.\n');
}

async function main() {
  const { Octokit } = await import('@octokit/rest');

  const octokit = new Octokit({ 
    auth: process.env.GITHUB_TOKEN,
    retry: { enabled: false }
  });

  // Create submodules directory if it doesn't exist
  await fs.mkdir(submodulesDir, { recursive: true });

  const maxRepos = process.env.MAX_REPOS ? parseInt(process.env.MAX_REPOS) : Infinity;
  let processedRepos = 0;

  try {
    for await (const response of octokit.paginate.iterator(octokit.repos.listForOrg, {
      org: githubOrg,
      per_page: 100
    })) {
      for (const repo of response.data) {

        try {
          await processRepository(repo, octokit);
        } catch (error) {
          console.error(`Error processing repository ${repo.name}:`, error);
          throw error; // Re-throw the error to stop the script
        }
        
        processedRepos++;
        if (processedRepos >= maxRepos) {
          console.log(`Reached the maximum number of repositories to process (${maxRepos})`);
          return;
        }
      }
    }
  } finally {
    rl.close();
    // Optionally, clean up submodules directory
    // await fs.rm(submodulesDir, { recursive: true, force: true });
  }
}

main().catch(error => {
  console.error("Script stopped due to an error:", error);
  process.exit(1);
});